feature: Support for Braket Direct Reservations (#87)

* [Feature]: Support for Braket Direct reservations

* [Feature]: Support for Braket Direct reservations - adding docstrings and tests

* [Feature]: Support for Braket Direct reservations - minor cleanup

* [Feature]: Support for Braket Direct reservations - Adding more tests

* [Feature]: Support for Braket Direct reservations - Final Polish

* [Feature]: Support for Braket Direct reservations - Improving Indentation

* [Feature]: Support for Braket Direct reservations - codereview

* [Feature]: Support for Braket Direct reservations - codereview

* [Feature]: Support for Braket Direct reservations - codereview

* [Feature]: Support for Braket Direct reservations - codereview suggestions

* [Feature]: Support for Braket Direct reservations - formatting

* [Feature]: Support for Braket Direct reservations - formatting

* Adding codereview suggestions

* Minor Formatting

* Minor Formatting

* adding direct reservation details in docstring

* Minor formatting and cleanup

* CI documentation fix

---------

Co-authored-by: Fe-r-oz <[email protected]>

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
+```

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
+```

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

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

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")