Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feature: Support for Braket Direct Reservations #87

Merged
merged 19 commits into from
Aug 28, 2024

Conversation

Fe-r-oz
Copy link

@Fe-r-oz Fe-r-oz commented Jun 8, 2024

Resolving Issue #80

Description of changes: Amazon Braket now has the Braket Direct feature, which allows users to reserve certain Direct-only devices or access experimental capabilities of devices. Implementation of Python SDK? https://github.com/amazon-braket/amazon-braket-sdk-python/blob/main/src/braket/aws/direct_reservations. This allows users to provide a reservation arn to access their reservation, as in the BDK.

I will add docstrings and more tests, later today!

Testing done:

julia> reservation = Braket.DirectReservation(DEVICE_ARN, RESERVATION_ARN)
               # Start reservation
Braket.DirectReservation("arn:aws:braket:us-east-1:123456789:device/qpu/ionq/Forte-1", "arn:aws:braket:us-east-1:123456789:reservation/uuid", false)

julia> Braket.start_reservation!(reservation)
"arn:aws:braket:us-east-1:123456789:reservation/uuid"

julia> @test reservation.is_active == true
Test Passed

julia> @test ENV["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] == DEVICE_ARN
Test Passed

julia> @test ENV["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] == RESERVATION_ARN
               # Stop reservation
Test Passed

julia> Braket.stop_reservation!(reservation)
Base.EnvDict with 62 entries:
  "SHELL"                     => "/bin/bash"
  "WINDOWID"                  => "0"
  "COLORTERM"                 => "truecolor"
  "XDG_SESSION_PATH"          => "/org/freedesktop/DisplayManager/Session0"
  "GNOME_DESKTOP_SESSION_ID"  => "this-is-deprecated"
  "GTK_IM_MODULE"             => "fcitx"
  "CONDA_EXE"                 => "/home/Transcendence/miniconda3/bin/conda"
  "_CE_M"                     => ""
  "LANGUAGE"                  => "en_US"
  "D_DISABLE_RT_SCREEN_SCALE" => "1"
  "SSH_AUTH_SOCK"             => "/run/user/1000/keyring/ssh"
  "XDG_DATA_HOME"             => "/home/Transcendence/.local/share"
  "INPUT_METHOD"              => "fcitx"
  "XDG_CONFIG_HOME"           => "/home/Transcendence/.config"
  "XMODIFIERS"                => "@im=fcitx"
  "DESKTOP_SESSION"           => "deepin"
  "XDG_SEAT"                  => "seat0"
  "PWD"                       => "/home/Transcendence"
  "XDG_SESSION_DESKTOP"       => "deepin"
  "LOGNAME"                   => "Transcendence"
  ⋮                           => ⋮

julia> @test reservation.is_active == false
Test Passed

julia> @test !haskey(ENV, "AMZN_BRAKET_RESERVATION_DEVICE_ARN")
Test Passed

julia> @test !haskey(ENV, "AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN")
Test Passed

julia> reservation = Braket.DirectReservation(DEVICE_ARN, RESERVATION_ARN)
Braket.DirectReservation("arn:aws:braket:us-east-1:123456789:device/qpu/ionq/Forte-1", "arn:aws:braket:us-east-1:123456789:reservation/uuid", false)

julia> @test reservation.device_arn == DEVICE_ARN
Test Passed

julia> @test reservation.reservation_arn == RESERVATION_ARN
Test Passed

julia> @test reservation.is_active == false
Test Passed

Note: I am new to git-secrets, so please bear with me!

Merge Checklist

Put an x in 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

  • I have added tests that prove my fix is effective or that my feature works (if appropriate)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

@Fe-r-oz
Copy link
Author

Fe-r-oz commented Jun 8, 2024

Also, Added this function to demonstrate that we can add actions or assertions after starting reservation and before stopping reservation! I hope this makes sense. In julia, instead of AbstractContextManager which is available in python, we can use do notation (https://groups.google.com/g/julia-users/c/leNMURKreZo) to achieve the same functionality.

function direct_reservation(reservation::DirectReservation, func::Function)
    env_vars = Dict(
        "AMZN_BRAKET_RESERVATION_DEVICE_ARN" => reservation.device_arn,
        "AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN" => reservation.reservation_arn
    )
    withenv(pairs(env_vars)...) do
        start_reservation!(reservation)
        try
            func()
        finally
            stop_reservation!(reservation)
        end
    end
end

    function test_func()
        println("Executing within reservation context")
	   return 5 
        # Add actions as needed
        @test ENV["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] == DEVICE_ARN
        @test ENV["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] == RESERVATION_ARN 
    end

    @testset "Direct Reservation Function" begin
	   reservation = Braket.DirectReservation(DEVICE_ARN, RESERVATION_ARN)
	   @test Braket.direct_reservation(reservation, test_func) == 5
    end

@rmshaffer rmshaffer linked an issue Jun 10, 2024 that may be closed by this pull request
src/device.jl Outdated Show resolved Hide resolved
src/device.jl Outdated Show resolved Hide resolved
src/device.jl Outdated Show resolved Hide resolved
src/device.jl Outdated Show resolved Hide resolved
test/device.jl Outdated Show resolved Hide resolved
@Fe-r-oz Fe-r-oz requested a review from kshyatt-aws June 13, 2024 02:07
@Fe-r-oz
Copy link
Author

Fe-r-oz commented Jun 13, 2024

Tested Locally:

Executing within reservation context
Test Summary:           | Pass  Total  Time
DirectReservation Tests |   11     11  0.3s
Test.DefaultTestSet("DirectReservation Tests", Any[Test.DefaultTestSet("Creating DirectReservation", Any[], 3, false, false, true, 1.718244271432349e9, 1.718244271581824e9, false, "REPL[26383]"), Test.DefaultTestSet("Starting and Stopping Reservation", Any[], 6, false, false, true, 1.71824427158189e9, 1.718244271607278e9, false, "REPL[26383]"), Test.DefaultTestSet("Direct Reservation Function", Any[], 1, false, false, true, 1.718244271607491e9, 1.71824427169584e9, false, "REPL[26383]"), Test.DefaultTestSet("Invalid Device Type", Any[], 1, false, false, true, 1.718244271695888e9, 1.71824427175638e9, false, "REPL[26383]")], 0, false, false, true, 1.718244271424501e9, 1.718244271756389e9, false, "REPL[26383]")

@Fe-r-oz Fe-r-oz changed the title [Feature]: Support for Braket Direct reservations feature: Support for Braket Direct Reservations Jun 14, 2024
src/device.jl Outdated
"""
DirectReservation(device::Union{Device, String, Nothing}, reservation_arn::Union{String, Nothing})

A context manager that adjusts AwsQuantumTasks generated within the context to utilize a reservation ARN
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why no doc link here for AwsQuantumTask?

src/device.jl Outdated
reserved device within the specified start and end times.

Arguments:
- device (Device | String | Nothing): The Braket device for which you possess a reservation ARN, or
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Device here can also use a doc link

src/device.jl Outdated
Arguments:
- device (Device | String | Nothing): The Braket device for which you possess a reservation ARN, or
alternatively, the device ARN.
- reservation_arn (String | Nothing): The Braket Direct reservation ARN to be implemented for all
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This doesn't match the actual struct field types

src/device.jl Outdated
# Start reservation function
function start_reservation!(state::DirectReservation)
if state.is_active
error("Another reservation is already active.")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not include the reservation_arn in this error message?

src/device.jl Outdated
try
func()
catch e
error("Error during reservation with device ARN $(reservation.device_arn): $(e)")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not show the reservation ARN in this error message as well?

Copy link
Member

@kshyatt-aws kshyatt-aws left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It might also make sense to create a new page in docs/src for reservations, since they're a pretty advanced feature that need some explanation. This could be a stub for now to be fleshed out later.

@kshyatt-aws
Copy link
Member

Other than a last few small changes this looks pretty ready to go to me.

@Fe-r-oz Fe-r-oz requested a review from kshyatt-aws June 19, 2024 04:38
@Fe-r-oz
Copy link
Author

Fe-r-oz commented Jun 19, 2024

Local Tests for Direct Reservations:


Precompiling project...
  1 dependency successfully precompiled in 21 seconds. 111 already precompiled.
     Testing Running tests...
Test Summary:    |Time
Method ambiguity | None  0.0s
Test Summary:           |Time
Unbound type parameters | None  0.0s
Test Summary:     | Pass  Total  Time
Undefined exports |    1      1  0.1s
Test Summary:                              | Pass  Total  Time
Compare Project.toml and test/Project.toml |    1      1  0.2s
Test Summary:      | Pass  Total  Time
Stale dependencies |    1      1  4.8s
Test Summary: | Pass  Total  Time
Compat bounds |    4      4  0.5s
Test Summary: |Time
Piracy        | None  0.0s
Test Summary:    |Time
Persistent tasks | None  0.0s
[ Info: Testing Braket-unit
WARNING: Method definition activated() in module Mocking at /home/Transcendence/.julia/packages/Mocking/U3Rvo/src/Mocking.jl:15 overwritten at /home/Transcendence/.julia/packages/Mocking/U3Rvo/src/Mocking.jl:24.
Test Summary: | Pass  Total     Time
AHS           | 1051   1051  1m12.4s
Test Summary: | Pass  Total  Time
utils         |   13     13  4.3s
Test Summary:      | Pass  Total  Time
QubitSet and Qubit |   14     14  0.4s
Test Summary:       | Pass  Total  Time
Dwave Device Schema |   15     15  2.8s
Test Summary:      | Pass  Total  Time
Ionq Device Schema |    4      4  1.2s
Test Summary:  | Pass  Total  Time
Rigetti device |    4      4  1.4s
Test Summary:               | Pass  Total  Time
Gate model simulator device |    3      3  0.2s
Test Summary:     | Pass  Total  Time
OQC device schema |    4      4  1.3s
Test Summary: | Pass  Total  Time
Quera device  |    2      2  0.7s
Test Summary: | Pass  Total  Time
Xanadu device |    3      3  0.8s
Test Summary:  | Pass  Total   Time
IR Translation |  132    132  27.6s
Test Summary: | Pass  Total     Time
Schemas       |  219    219  4m13.1s
Test Summary: | Pass  Total  Time
Devices       |   54     54  6.6s
Test Summary: | Pass  Total   Time
Circuit       |  284    284  43.2s
Test Summary:    | Pass  Total  Time
Measure operator |    7      7  0.0s
Test Summary:   | Pass  Total  Time
Free parameters |    6      6  1.2s
Test Summary: | Pass  Total   Time
Gates         |  349    349  14.3s
Test Summary: | Pass  Total  Time
Observables   |  250    250  6.4s
Test Summary: | Pass  Total  Time
Noises        |   62     62  2.3s
Test Summary: | Pass  Total  Time
Noise models  |   90     90  4.3s
Test Summary:              | Pass  Total  Time
Compiler directives basics |    8      8  0.1s
Test Summary:              | Pass  Total   Time
GateModelQuantumTaskResult |   11     11  18.4s
Test Summary:               | Pass  Total  Time
shots>0 results computation |   11     11  4.0s
Test Summary:             | Pass  Total  Time
PhotonicQuantumTaskResult |    5      5  6.2s
Test Summary:              | Pass  Broken  Total   Time
AnnealingQuantumTaskResult |   13       1     14  20.3s
Test Summary: | Pass  Total  Time
Cost tracker  |   15     15  1.7s
Test Summary: | Pass  Total   Time
Tasks         |  128    128  10.8s
Test Summary: | Pass  Total  Time
Batched tasks |   42     42  4.5s
Test Summary: | Pass  Total   Time
Local Jobs    |  101    101  10.0s
Test Summary: | Pass  Total  Time
Job macro     |    4      4  5.7s
Test Summary: | Pass  Total     Time
Jobs          |   55     55  1m44.4s
     Testing Braket tests passed 

@@ -0,0 +1,8 @@
# Reservations

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should probably have some flavour text here explaining what reservations are and why anyone would use them.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please checkout the documentation.

@Fe-r-oz Fe-r-oz requested a review from kshyatt-aws June 19, 2024 20:35
Copy link
Member

@kshyatt-aws kshyatt-aws left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can be merged into a feature branch (should add more integration tests that are difficult for UH participants to do themselves)

@kshyatt-aws kshyatt-aws changed the base branch from main to ksh/feature_bump August 28, 2024 17:06
@kshyatt-aws kshyatt-aws merged commit fa9c25e into amazon-braket:ksh/feature_bump Aug 28, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Support for Braket Direct reservations
2 participants