feature: Support for Braket Direct Reservations

src/device.jl

@@ -336,3 +336,70 @@ 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::Union{Device, String, Nothing}, reservation_arn::Union{String, Nothing})
+
+A context manager that adjusts AwsQuantumTasks 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 (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
+  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
+function start_reservation!(state::DirectReservation)
+    if state.is_active
+        error("Another reservation is already active.")
+    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
+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
+
+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): $(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")