fix: Bump dep versions and fix various bugs (#24)

- Fix errors in cost tracker.
- `AwsQuantumTask` now tracks the `AWSConfig` it was run with (inc.
  region).
- Bump dependency versions
- Fix `test/runtests.jl` to not develop `PyBraket.jl` until necessary.
- Add a CI variable to turn off Aqua tests that break when running in CI.
- Restore non-AG integration tests for `Braket.jl`.
- Allow optional arguments in schema for device parsing.

.github/workflows/CI.yml

@@ -28,7 +28,7 @@ jobs:
         group:
           - Braket-unit
           - PyBraket-unit
-        os: [windows-latest, macOS-latest]        
+        os: [windows-latest, macOS-latest]
         arch: ['x64']
     steps:
       - uses: actions/checkout@v1
@@ -41,10 +41,20 @@ jobs:
       - run: |
           git config --global user.name Tester
           git config --global user.email [email protected]
+      - name: "Dev PyBraket package on Windows"
+        if: ${{ matrix.os == 'windows-latest' }}
+        run: |
+          julia --project -e 'using Pkg; Pkg.develop(path=joinpath(pwd(), \"PyBraket\"))'
+          # must escape for windows
+      - name: "Dev PyBraket package on OSX"
+        if: ${{ matrix.os == 'macOS-latest' }}
+        run: |
+          julia --project -e 'using Pkg; Pkg.develop(path=joinpath(pwd(), "PyBraket"))'
       - name: "Run tests"
         uses: julia-actions/julia-runtest@v1
         env:
           GROUP: ${{ matrix.group }}
+          BRAKET_CI: true
           JULIA_CONDAPKG_VERBOSITY: 2
   version-test:
     runs-on: ubuntu-latest 
@@ -72,10 +82,14 @@ jobs:
       - run: |
           git config --global user.name Tester
           git config --global user.email [email protected]
+      - name: "Dev PyBraket package"
+        run: |
+          julia --project -e 'using Pkg; Pkg.develop(path=joinpath(pwd(), "PyBraket"))'
       - name: "Run tests"
         uses: julia-actions/julia-runtest@v1
         env:
           JULIA_CONDAPKG_VERBOSITY: 2
+          BRAKET_CI: true
           GROUP: ${{ matrix.group }}
       - name: "Process coverage"
         uses: julia-actions/julia-processcoverage@v1

Project.toml

@@ -30,26 +30,28 @@ Tar = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
-AWS = "=1.78.0"
+AWS = "=1.81.0"
 AWSS3 = "=0.10.2"
-Aqua = "0.5"
+Aqua = "=0.6"
 AxisArrays = "=0.4.6"
-CSV = "=0.10.4"
-Compat = "=4.3.0"
+CSV = "=0.10.9"
+Compat = "=4.5.0"
 DataStructures = "=0.18.13"
 DecFP = "=1.3.1"
 Distributions = "=0.25.76"
 Graphs = "=1.7.4"
-HTTP = "=1.4.0"
-JSON3 = "=1.11.1"
-Mocking = "=0.7.2"
-NamedTupleTools = "=0.14.1"
+HTTP = "=1.7.3"
+JSON3 = "=1.12.0"
+Mocking = "=0.7.5"
+NamedTupleTools = "=0.14.3"
 OrderedCollections = "=1.4.1"
 StructTypes = "=1.10.0"
 Tar = "1.9.3"
 julia = "1.6"
 
 [extras]
+AWS = "fbe9abb3-538b-5e4e-ba9e-bc94f4f92ebc"
+AWSS3 = "1c724243-ef5b-51ab-93f4-b0a88ac62a95"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DecFP = "55939f99-70c6-5e9b-8bb0-5071ed7d61fd"

PyBraket/Project.toml

@@ -4,7 +4,6 @@ authors = ["Katharine Hyatt <[email protected]>"]
 version = "0.2.0"
 
 [deps]
-AWS = "fbe9abb3-538b-5e4e-ba9e-bc94f4f92ebc"
 Braket = "19504a0f-b47d-4348-9127-acc6cc69ef67"
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -14,18 +13,17 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructTypes = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
 
 [compat]
-AWS = "=1.78.0"
-Aqua = "0.5"
-Braket = "=0.2.0"
-CondaPkg = "=0.2.14"
+Aqua = "=0.6"
+Braket = "0.2.0"
+CondaPkg = "=0.2.15"
 DataStructures = "=0.18.13"
-PythonCall = "=0.9.7"
+PythonCall = "=0.9.10"
 StructTypes = "=1.10.0"
 julia = "1.6"
 
 [extras]
-AWS = "fbe9abb3-538b-5e4e-ba9e-bc94f4f92ebc"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Braket = "19504a0f-b47d-4348-9127-acc6cc69ef67"
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"

PyBraket/test/Project.toml

@@ -1,5 +1,4 @@
 [deps]
-AWS = "fbe9abb3-538b-5e4e-ba9e-bc94f4f92ebc"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Braket = "19504a0f-b47d-4348-9127-acc6cc69ef67"
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
@@ -8,10 +7,3 @@ PyBraket = "e85266a6-1825-490b-a80e-9b9469c53660"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[compat]
-AWS = "=1.78.0"
-Aqua = "0.5"
-Braket = "=0.2.0"
-CondaPkg = "=0.2.14"
-PythonCall = "=0.9.7"

PyBraket/test/runtests.jl

@@ -1,16 +1,16 @@
-using Test, Aqua, AWS, PyBraket
+using Test, Aqua, Braket, Braket.AWS, PyBraket
 
-Aqua.test_all(PyBraket, ambiguities=false, unbound_args=false)
+Aqua.test_all(PyBraket, ambiguities=false, unbound_args=false, piracy=false)
 Aqua.test_ambiguities(PyBraket)
 
 function set_aws_creds(test_type)
     if test_type == "unit"
-        creds = AWS.AWSCredentials("", "")
-        config = AWS.AWSConfig(creds, "", "")
-        AWS.global_aws_config(config)
+        creds = Braket.AWS.AWSCredentials("", "")
+        config = Braket.AWS.AWSConfig(creds, "", "")
+        Braket.AWS.global_aws_config(config)
     elseif test_type == "integ"
         # should pickup correct creds from envvars 
-        AWS.global_aws_config()
+        Braket.AWS.global_aws_config()
     else
         throw(ArgumentError("invalid test_type $test_type, must be one of 'integ' or 'unit'"))
     end

src/device.jl

@@ -35,13 +35,13 @@ Base.convert(::Type{String}, d::AwsDevice) = d._arn
 Base.show(io::IO, d::AwsDevice) = print(io, "AwsDevice(arn="*d._arn*")")
 function (d::AwsDevice)(task_spec::Union{Circuit, AnalogHamiltonianSimulation, AbstractProgram}; s3_destination_folder=default_task_bucket(), shots=nothing, poll_timeout_seconds::Int=DEFAULT_RESULTS_POLL_TIMEOUT, poll_interval_seconds::Int=DEFAULT_RESULTS_POLL_INTERVAL, inputs=Dict{String, Float64}(), kwargs...)
     shots_ = isnothing(shots) ? d._default_shots : shots
-    return AwsQuantumTask(d._arn, task_spec, s3_destination_folder=s3_destination_folder, shots=shots_, poll_timeout_seconds=poll_timeout_seconds, poll_interval_seconds=poll_interval_seconds, inputs=inputs, kwargs...)
+    return AwsQuantumTask(d._arn, task_spec, s3_destination_folder=s3_destination_folder, shots=shots_, poll_timeout_seconds=poll_timeout_seconds, poll_interval_seconds=poll_interval_seconds, inputs=inputs, config=d._config, kwargs...)
 end
 
 # currently no batch support for AHS
-function (d::AwsDevice)(task_specs::Vector{<:Union{Circuit, AbstractProgram}}; s3_destination_folder=default_task_bucket(), shots=nothing, max_parallel=nothing, poll_timeout_seconds::Int=DEFAULT_RESULTS_POLL_TIMEOUT, poll_interval_seconds::Int=DEFAULT_RESULTS_POLL_INTERVAL, inputs=Dict{String, Float64}(), kwargs...)
+function (d::AwsDevice)(task_specs::Vector{<:Union{Circuit, AbstractProgram}}; s3_destination_folder=default_task_bucket(), shots=nothing, max_parallel=nothing, poll_timeout_seconds::Int=DEFAULT_RESULTS_POLL_TIMEOUT, poll_interval_seconds::Int=DEFAULT_RESULTS_POLL_INTERVAL, inputs=Dict{String, Float64}(), config=d._config, kwargs...)
     shots_ = isnothing(shots) ? d._default_shots : shots
-    return AwsQuantumTaskBatch(d._arn, task_specs; s3_destination_folder=s3_destination_folder, shots=shots_, poll_timeout_seconds=poll_timeout_seconds, poll_interval_seconds=poll_interval_seconds, inputs=inputs, kwargs...)
+    return AwsQuantumTaskBatch(d._arn, task_specs; s3_destination_folder=s3_destination_folder, shots=shots_, poll_timeout_seconds=poll_timeout_seconds, poll_interval_seconds=poll_interval_seconds, inputs=inputs, config=config, kwargs...)
 end
 
 function _construct_topology_graph(d::AwsDevice)
@@ -222,6 +222,7 @@ function get_devices(; arns::Vector{String}=String[], names::Vector{String}=Stri
         config_for_region = region == current_region ? global_config : AWSConfig(global_config.credentials, region, global_config.output)
         # Simulators are only instantiated in the same region as the AWS session
         types_for_region = string.(sort(region == current_region ? types : setdiff(types, "SIMULATOR")))
+        types_for_region = isnothing(types_for_region) ? Vector[] : types_for_region
         region_devices = search_devices(arns=arns, names=names,
                                         types=types_for_region,
                                         statuses=statuses,

src/raw_schema.jl

@@ -996,7 +996,7 @@ StructTypes.StructType(::Type{BlackbirdDeviceActionProperties}) = StructTypes.Un
 struct JaqcdDeviceActionProperties <: DeviceActionProperties
     version::Vector{String}
     actionType::String
-    supportedOperations::Vector{String}
+    supportedOperations::Union{Nothing, Vector{String}}
     supportedResultTypes::Union{Nothing, Vector{ResultType}}
     disabledQubitRewiringSupported::Union{Nothing, Bool}
 end
@@ -1005,7 +1005,7 @@ StructTypes.StructType(::Type{JaqcdDeviceActionProperties}) = StructTypes.Unorde
 struct OpenQASMDeviceActionProperties <: DeviceActionProperties
     version::Vector{String}
     actionType::String
-    supportedOperations::Vector{String}
+    supportedOperations::Union{Nothing, Vector{String}}
     supportedPragmas::Vector{String}
     forbiddenPragmas::Vector{String}
     maximumQubitArrays::Union{Nothing, Int}

src/task.jl

@@ -16,15 +16,17 @@ mutable struct AwsQuantumTask
     poll_timeout_seconds::Int
     poll_interval_seconds::Int
     _future::Union{Future, Nothing}
+    _config::AbstractAWSConfig
     _metadata::Dict{String, Any}
     _logger::AbstractLogger
     _result::Union{Nothing, AbstractQuantumTaskResult}
     function AwsQuantumTask(arn::String;
                    client_token::String=string(uuid1()),
                    poll_timeout_seconds::Int=DEFAULT_RESULTS_POLL_TIMEOUT,
                    poll_interval_seconds::Int=DEFAULT_RESULTS_POLL_INTERVAL,
-                   logger=global_logger())
-        new(arn, client_token, poll_timeout_seconds, poll_interval_seconds, nothing, Dict(), logger, nothing)
+                   logger=global_logger(),
+                   config::AWSConfig=global_aws_config())
+        new(arn, client_token, poll_timeout_seconds, poll_interval_seconds, nothing, config, Dict(), logger, nothing)
     end
 end
 Base.show(io::IO, t::AwsQuantumTask) = println(io, "AwsQuantumTask(\"id/taskArn\":\"$(arn(t))\")")
@@ -45,15 +47,16 @@ function AwsQuantumTask(args::NamedTuple)
     s3_key_prefix = args[:outputS3KeyPrefix]
     shots        = args[:shots]
     extra_opts   = args[:extra_opts]
+    config       = args[:config]
     job_token    = get(ENV, "AMZN_BRAKET_JOB_TOKEN", nothing)
     !isnothing(job_token) && merge!(extra_opts, Dict("jobToken"=>job_token))
     timeout_seconds  = get(args, :poll_timeout_seconds, DEFAULT_RESULTS_POLL_TIMEOUT)
     interval_seconds = get(args, :poll_interval_seconds, DEFAULT_RESULTS_POLL_INTERVAL)
     merge!(AWS.AWSServices.braket.service_specific_headers, AWS.LittleDict("Braket-Trackers"=>string(length(GlobalTrackerContext[]))))
-    response         = BRAKET.create_quantum_task(action, client_token, device_arn, s3_bucket, s3_key_prefix, shots, extra_opts)
+    response         = BRAKET.create_quantum_task(action, client_token, device_arn, s3_bucket, s3_key_prefix, shots, extra_opts, aws_config=config)
     pop!(AWS.AWSServices.braket.service_specific_headers, "Braket-Trackers")
     broadcast_event!(TaskCreationEvent(response["quantumTaskArn"], shots, !isnothing(job_token), device_arn))
-    return AwsQuantumTask(response["quantumTaskArn"]; client_token=client_token, poll_timeout_seconds=timeout_seconds, poll_interval_seconds=interval_seconds)
+    return AwsQuantumTask(response["quantumTaskArn"]; client_token=client_token, poll_timeout_seconds=timeout_seconds, poll_interval_seconds=interval_seconds, config=config)
 end
 
 function default_task_bucket()
@@ -94,8 +97,9 @@ function AwsQuantumTask(device_arn::String,
                         poll_timeout_seconds::Int=DEFAULT_RESULTS_POLL_TIMEOUT,
                         poll_interval_seconds::Int=DEFAULT_RESULTS_POLL_INTERVAL,
                         tags::Dict{String, String}=Dict{String, String}(),
-                        inputs::Dict{String,Float64}=Dict{String, Float64}())
-    args = prepare_task_input(task_spec, device_arn, s3_destination_folder, shots, device_params, disable_qubit_rewiring, tags=tags, poll_timeout_seconds=poll_timeout_seconds, poll_interval_seconds=poll_interval_seconds, inputs=inputs)
+                        inputs::Dict{String,Float64}=Dict{String, Float64}(),
+                        config::AbstractAWSConfig=global_aws_config())
+    args = prepare_task_input(task_spec, device_arn, s3_destination_folder, shots, device_params, disable_qubit_rewiring, tags=tags, poll_timeout_seconds=poll_timeout_seconds, poll_interval_seconds=poll_interval_seconds, inputs=inputs, config=config)
     return AwsQuantumTask(args)
 end
 
@@ -123,7 +127,8 @@ _create_annealing_device_params(device_params, device_arn) = _create_annealing_d
 function _create_common_params(device_arn::String, s3_destination_folder::Tuple{String, String}, shots::Int; kwargs...)
     timeout_seconds  = get(kwargs, :poll_timeout_seconds, DEFAULT_RESULTS_POLL_TIMEOUT)
     interval_seconds = get(kwargs, :poll_interval_seconds, DEFAULT_RESULTS_POLL_INTERVAL)
-    return (device_arn=device_arn, outputS3Bucket=s3_destination_folder[1],  outputS3KeyPrefix=s3_destination_folder[2], shots=shots, poll_timeout_seconds=timeout_seconds, poll_interval_seconds=interval_seconds)
+    config           = get(kwargs, :config, global_aws_config())
+    return (device_arn=device_arn, outputS3Bucket=s3_destination_folder[1],  outputS3KeyPrefix=s3_destination_folder[2], shots=shots, poll_timeout_seconds=timeout_seconds, poll_interval_seconds=interval_seconds, config=config)
 end
 
 function prepare_task_input(problem::Problem, device_arn::String, s3_folder::Tuple{String, String}, shots::Int, device_params::Union{Dict{String, Any}, DwaveDeviceParameters, DwaveAdvantageDeviceParameters, Dwave2000QDeviceParameters}, disable_qubit_rewiring::Bool=false; kwargs...)
@@ -231,7 +236,7 @@ Cancels the task `t`.
 """
 function cancel(t::AwsQuantumTask)
     #!isnothing(t._future) && cancel(t.future)
-    resp = BRAKET.cancel_quantum_task(t.client_token, t.arn)
+    resp = BRAKET.cancel_quantum_task(t.client_token, HTTP.escapeuri(t.arn), aws_config=t._config)
     broadcast_event!(TaskStatusEvent(t.arn, resp["cancellationStatus"]))
     return
 end

src/task_batch.jl

@@ -16,9 +16,10 @@ mutable struct AwsQuantumTaskBatch
     _shots::Int
     _poll_timeout_seconds::Int
     _poll_interval_seconds::Int
-    function AwsQuantumTaskBatch(tasks, results, unsuccessful, device_arn, task_specs, s3_dest, shots, timeout, interval)
+    _config::AbstractAWSConfig
+    function AwsQuantumTaskBatch(tasks, results, unsuccessful, device_arn, task_specs, s3_dest, shots, timeout, interval, config=global_aws_config())
         length(tasks) != length(task_specs) && throw(ArgumentError("number of quantum tasks ($(length(tasks))) and task specifications ($(length(task_specs))) must be equal!"))
-        new(tasks, results, unsuccessful, device_arn, task_specs, s3_dest, shots, timeout, interval)
+        new(tasks, results, unsuccessful, device_arn, task_specs, s3_dest, shots, timeout, interval, config)
     end
     @doc """
         AwsQuantumTaskBatch(device_arn::String, task_specs::Vector{<:Union{AbstractProgram, Circuit}}; kwargs...) -> AwsQuantumTaskBatch
@@ -31,9 +32,9 @@ mutable struct AwsQuantumTaskBatch
       - `poll_timeout_seconds::Int` - maximum number of seconds to wait while polling for results. Default: $DEFAULT_RESULTS_POLL_TIMEOUT
       - `poll_interval_seconds::Int` - default number of seconds to wait between attempts while polling for results. Default: $DEFAULT_RESULTS_POLL_INTERVAL
     """
-    function AwsQuantumTaskBatch(device_arn::String, task_specs::Vector{<:Union{AbstractProgram, Circuit}}; s3_destination_folder::Tuple{String, String}=default_task_bucket(), shots::Int=DEFAULT_SHOTS, poll_timeout_seconds::Int=DEFAULT_RESULTS_POLL_TIMEOUT, poll_interval_seconds=DEFAULT_RESULTS_POLL_INTERVAL, inputs=Dict{String, Float64}())
-        tasks = launch_batch(device_arn, task_specs; s3_destination_folder=s3_destination_folder, shots=shots, poll_interval_seconds=poll_interval_seconds, inputs=inputs)
-        new(tasks, nothing, Set(), device_arn, task_specs, s3_destination_folder, shots, poll_timeout_seconds, poll_interval_seconds)
+    function AwsQuantumTaskBatch(device_arn::String, task_specs::Vector{<:Union{AbstractProgram, Circuit}}; s3_destination_folder::Tuple{String, String}=default_task_bucket(), shots::Int=DEFAULT_SHOTS, poll_timeout_seconds::Int=DEFAULT_RESULTS_POLL_TIMEOUT, poll_interval_seconds=DEFAULT_RESULTS_POLL_INTERVAL, inputs=Dict{String, Float64}(), config::AbstractAWSConfig=global_aws_config())
+        tasks = launch_batch(device_arn, task_specs; s3_destination_folder=s3_destination_folder, shots=shots, poll_interval_seconds=poll_interval_seconds, inputs=inputs, config=config)
+        new(tasks, nothing, Set(), device_arn, task_specs, s3_destination_folder, shots, poll_timeout_seconds, poll_interval_seconds, config)
     end
 end
 

src/tracker.jl

@@ -26,7 +26,7 @@ end
 function price_search(p::Pricing; kwargs...)
     isempty(p._price_list) && get_prices!(p)
     return filter(p._price_list) do entry
-        all(haskey(entry, string(k)) && entry[string(k)] == v for (k, v) in kwargs)
+        all([haskey(entry, string(k)) && !ismissing(entry[string(k)]) && entry[string(k)] == v for (k, v) in kwargs])
     end
 end
 price_search(; kwargs...) = price_search(Prices[]; kwargs...)
@@ -128,14 +128,17 @@ function _get_qpu_task_cost(arn::String, details::Dict{String})
     occursin("2000Q", device_name) && (device_name = "2000Q")
     occursin("Advantage_system", device_name) && (device_name = "Advantage_system")
 
-    search_dict[Symbol("Device name")] = device_name
     if details["job_task"]
+        search_dict[Symbol("Device Name")] = device_name
         shot_product_family = "Braket Managed Jobs QPU Task Shot"
         task_product_family = "Braket Managed Jobs QPU Task"
     else
+        search_dict[Symbol("DeviceName")] = device_name
         shot_product_family = "Quantum Task-Shot"
         task_product_family = "Quantum Task"
     end
+    shot_prices = price_search(; namedtuple(search_dict)...)
+
     search_dict[Symbol("Product Family")] = shot_product_family
     shot_prices = price_search(; namedtuple(search_dict)...)
     search_dict[Symbol("Product Family")] = task_product_family
@@ -162,7 +165,7 @@ function _get_simulator_task_cost(arn::String, details::Dict{String})
 
     device_name = uppercase(String(split(details["device"], '/')[end]))
 
-    search_dict= Dict(Symbol("Device name")=>device_name)
+    search_dict = Dict(Symbol("Device name")=>device_name)
     if details["job_task"]
         product_family = "Braket Managed Jobs Simulator Task"
         operation = "Managed-Jobs"

test/Project.toml

@@ -1,4 +1,6 @@
 [deps]
+AWS = "fbe9abb3-538b-5e4e-ba9e-bc94f4f92ebc"
+AWSS3 = "1c724243-ef5b-51ab-93f4-b0a88ac62a95"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Braket = "19504a0f-b47d-4348-9127-acc6cc69ef67"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
@@ -10,6 +12,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Mocking = "78c3b35d-d492-501b-9361-3d52fe80e533"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+PyBraket = "e85266a6-1825-490b-a80e-9b9469c53660"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructTypes = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"