diff --git a/docs/src/device.md b/docs/src/device.md index f1bab3a6..052635e9 100644 --- a/docs/src/device.md +++ b/docs/src/device.md @@ -7,4 +7,8 @@ Braket.BraketDevice isavailable search_devices get_devices -``` +Braket.DirectReservation +Braket.start_reservation! +Braket.stop_reservation! +Braket.direct_reservation +``` \ No newline at end of file diff --git a/docs/src/reservations.md b/docs/src/reservations.md new file mode 100644 index 00000000..733c4b07 --- /dev/null +++ b/docs/src/reservations.md @@ -0,0 +1,27 @@ +# Reservations + +[Reservations](https://docs.aws.amazon.com/braket/latest/developerguide/braket-reservations.html) grant exclusive access to a specific quantum device for a predetermined time slot. This control over execution windows offers several advantages to users: + + * **Predictability:** Users have guaranteed knowledge of precisely when their tasks will be executed, eliminating scheduling uncertainties. + * **Prioritization:** During the reservation window, the user's workloads take precedence over others, avoiding queueing delays and potential bottlenecks. + * **Efficiency:** The cost is based solely on the reserved duration, irrespective of the number of tasks the user runs within that window. + +# Why Use Reservations? + +In certain scenarios, reservations provide significant benefits over on-demand access: + +- **Production Runs**: When finalizing research or performing critical computations, reservations guarantee timely completion, ensuring adherence to deadlines. +- **Live Demos or Workshops**: Reservations secure exclusive device access for showcasing work or conducting workshops at specific times. +- **Streamlined Workflows**: Users can schedule reservations for tasks requiring execution at particular moments within their workflows, optimizing their overall process. + +# Reservations vs. On-Demand Access + +On-demand access is better suited for the initial development and prototyping stages of a quantum project. This method allows for rapid iterations without the need for pre-scheduled reservations, facilitating agile development. However, once the project progresses towards final production runs requiring guaranteed execution times, reservations become the preferred choice. +""" + +```@docs +Braket.DirectReservation +Braket.start_reservation! +Braket.stop_reservation! +Braket.direct_reservation +``` \ No newline at end of file diff --git a/src/device.jl b/src/device.jl index 639134f9..4b59620e 100644 --- a/src/device.jl +++ b/src/device.jl @@ -340,3 +340,92 @@ function get_devices(; arns::Vector{String}=String[], names::Vector{String}=Stri end return sort(collect(values(device_map)), by=(x->getproperty(x, Symbol("_"*order_by)))) end + +""" + DirectReservation(device_arn::String, reservation_arn::String) + +A context manager that adjusts [`AwsQuantumTask`](@ref) generated within the context to utilize a +reservation ARN for all tasks targeting the designated device. Notably, this manager permits only +one reservation at a time. + +[Reservations](https://docs.aws.amazon.com/braket/latest/developerguide/braket-reservations.html) +are specific to AWS accounts and devices. Only the AWS account that initiated the reservation +can utilize the reservation ARN. Moreover, the reservation ARN is solely valid on the +reserved device within the specified start and end times. + +Arguments: +- device_arn: The [`BraketDevice`](@ref) for which you possess a reservation ARN, or + alternatively, the [`Device`](@ref) ARN. +- reservation_arn: The Braket [`DirectReservation`](@ref) ARN to be implemented for all + quantum tasks executed within the context. +""" +mutable struct DirectReservation + device_arn::String + reservation_arn::String + is_active::Bool +end + +DirectReservation(device_arn::String, reservation_arn::String) = DirectReservation(device_arn, reservation_arn, false) + +# Start reservation function +""" + start_reservation!(state::DirectReservation) + +[`start_reservation!`](@ref) activates a [`DirectReservation`](@ref) object, granting exclusive access to the +specified quantum device for the reserved time window. +""" +function start_reservation!(state::DirectReservation) + if state.is_active + error("Another reservation is already active. Reservation ARN: $reservation_arn") + end + state.is_active = true + ENV["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] = state.device_arn + ENV["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] = state.reservation_arn +end + +# Stop reservation function +""" + stop_reservation!(state::DirectReservation) + +[`stop_reservation!`](@ref) deactivates a [`DirectReservation`](@ref) object, releasing exclusive access to the quantum device. +""" +function stop_reservation!(state::DirectReservation) + if !state.is_active + @warn "Reservation is not active." + end + state.is_active = false + if haskey(ENV, "AMZN_BRAKET_RESERVATION_DEVICE_ARN") + delete!(ENV, "AMZN_BRAKET_RESERVATION_DEVICE_ARN") + end + if haskey(ENV, "AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN") + delete!(ENV, "AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN") + end +end + +""" + direct_reservation(reservation::DirectReservation, func::Function) + +Runs a function within the context of an AWS Braket [`DirectReservation`](@ref). This function sets +the necessary environment variables, starts the reservation, executes the provided function, +and then stops the reservation, ensuring proper resource management. + +Arguments: +- reservation::DirectReservation: The reservation details, including device ARN and reservation ARN. +- func::Function: The function to run within the reservation context. +""" +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() + catch e + error("Error during reservation with device ARN $(reservation.device_arn) and reservation ARN $(reservation.reservation_arn): $(e)") + finally + stop_reservation!(reservation) + end + end +end diff --git a/test/integ_tests/direct_reservations.jl b/test/integ_tests/direct_reservations.jl new file mode 100644 index 00000000..db8567b7 --- /dev/null +++ b/test/integ_tests/direct_reservations.jl @@ -0,0 +1,54 @@ +using Braket, Test, Mocking, JSON3, Dates, Graphs +using Braket: status +Mocking.activate() + +# Constants +RESERVATION_ARN = "arn:aws:braket:us-east-1:123456789:reservation/uuid" +DEVICE_ARN = "arn:aws:braket:us-east-1:123456789:device/qpu/ionq/Forte-1" + +@testset "DirectReservation Tests" begin + # Creating a DirectReservation + @testset "Creating DirectReservation" begin + reservation = Braket.DirectReservation(DEVICE_ARN, RESERVATION_ARN) + @test reservation.device_arn == DEVICE_ARN + @test reservation.reservation_arn == RESERVATION_ARN + @test reservation.is_active == false + end + + # Starting and stopping a reservation + @testset "Starting and Stopping Reservation" begin + reservation = Braket.DirectReservation(DEVICE_ARN, RESERVATION_ARN) + + # Start reservation + Braket.start_reservation!(reservation) + @test reservation.is_active == true + @test ENV["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] == DEVICE_ARN + @test ENV["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] == RESERVATION_ARN + + # Stop reservation + Braket.stop_reservation!(reservation) + @test reservation.is_active == false + @test !haskey(ENV, "AMZN_BRAKET_RESERVATION_DEVICE_ARN") + @test !haskey(ENV, "AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN") + + 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 + + @testset "Invalid Device Type" begin + invalid_device = 12345 # Not a valid device type + @test_throws UndefVarError DirectReservation(invalid_device, RESERVATION_ARN) + end + +end diff --git a/test/integ_tests/runtests.jl b/test/integ_tests/runtests.jl index 233e718d..9023de91 100644 --- a/test/integ_tests/runtests.jl +++ b/test/integ_tests/runtests.jl @@ -6,6 +6,7 @@ s3_destination_folder = Braket.default_task_bucket() include("adjoint_gradient.jl") include("create_local_quantum_job.jl") include("create_quantum_job.jl") +include("direct_reservations.jl") include("job_macro.jl") include("measure.jl") include("cost_tracking.jl")