Support nested decorators with Braket hybrid job and CUDA-Q kernel

[yitchen-tim]

Describe the feature you'd like
Support the follow use case that (1) Braket job is submitted with the @hybrid_job decorator, and (2) CUDA-Q kernel is created with @cudaq.kernel decorator. Currently, this use case throws a source code error.

from braket.jobs import hybrid_job

@hybrid_job(device="arn:aws:braket:::device/quantum-simulator/amazon/sv1")
def bell():
    import cudaq
    @cudaq.kernel
    def kernel():
        qubits = cudaq.qvector(2)
        h(qubits[0])
        cx(qubits[0], qubits[1])
    return cudaq.sample(kernel)

What is the root cause?
The error comes from the fact that cloudpickle, on which the @hybrid_job decorator is built, does not preserve the source code of functions that are defined in the hybrid_job decorated function (e.g., the source code of kernel() in the snippet above is not preserved). However, @cudaq.kernel relies on source code analysis to build a kernel behind the scene. A related issue is documented here.

What are the alternatives that works today?
There are three alternatives:

1. use AutoQASM

AutoQASM has a hybrid job decorator which behaves the same as Braket's hybrid job decorator except that it copies the source code of the inner functions. AutoQASM also modifies the code attribute of these functions such that their source code path points to the copied files. In this way, cudaq.kernel will be able to find the source code for kernel building.

You can simply replace the import with from autoqasm import hybrid_job to use this solution. However, this only works for function defined inside the hybrid_job decorated function (e.g., the kernel() will work). If the functions is defined in a different file, you may still see the same error.

2. Use non-decorator kernel creation to avoid source code analysis

If you use other way to program CUDA-Q, the kernel would be created without the need of source code analysis and you won't face the source code error. The code snippet below is an example.

from braket.jobs import hybrid_job

@hybrid_job(device="arn:aws:braket:::device/quantum-simulator/amazon/sv1")
def bell():
    import cudaq

    kernel = cudaq.make_kernel()
    qubits = kernel.qalloc(2)
    kernel.h(qubits[0])
    kernel.cx(qubits[0], qubits[1])
    return cudaq.sample(kernel)

3. Use non-decorator Braket job submission to avoid cloudpickle

If you create Braket job with the traditional interface, cloudpickle is not needed and you won't see the source code error. To do so, you first prepare a algorithm that includes your CUDA-Q program.

# content of example_algorithm_script.py
import cudaq

@cudaq.kernel
def kernel():
    qubits = cudaq.qvector(2)
    h(qubits[0])
    cx(qubits[0], qubits[1])

result =  cudaq.sample(kernel)
save_result(result) # a custom save result function

And in a Jupyter notebook or in command line, you submit the job with the algorithm script.

from braket.aws import AwsQuantumJob

job = AwsQuantumJob.create(
    device="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
    source_module="example_algorithm_script.py",
)