ローカルコードをハイブリッドジョブとして実行する

Amazon Braket Hybrid Jobs は、ハイブリッド量子古典アルゴリズムのフルマネージドオーケストレーションを提供し、Amazon EC2 コンピューティングリソースと Amazon Braket 量子処理ユニット (QPU) アクセスを統合します。ハイブリッドジョブで作成された量子タスクは、個々の量子タスクよりも優先的にキューイングされるため、量子タスクキュー内の変動によってアルゴリズムが中断されることはありません。各 QPU は独自のハイブリッドジョブキューを維持することで、ハイブリッドジョブを一度に 1 つだけ実行するようにします。

ローカル Python コードからハイブリッドジョブを作成する

ローカル Python コードを Amazon Braket Hybrid Jobs として実行できます。これは、次のコードサンプルに示すように、コードに @hybrid_job デコレータで注釈を付けることで実現できます。カスタム環境では、Amazon Elastic Container Registry (ECR) のカスタムコンテナの使用を選択できます。

@hybrid_job デコレータを使用して関数に注釈を付けることができます。Braket は、デコレータ内のコードを Braket ハイブリッドジョブのアルゴリズムスクリプトに変換します。次に、ハイブリッドジョブは Amazon EC2 インスタンス上で被デコレート関数を呼び出します。ジョブの進行状況は、job.state() または Braket コンソールでモニタリングできます。次のコード例は、State Vector Simulator (SV1) device で 5 つの状態のシーケンスを実行する方法を示しています。

from braket.aws import AwsDevice
from braket.circuits import Circuit, FreeParameter, Observable
from braket.devices import Devices
from braket.jobs.hybrid_job import hybrid_job
from braket.jobs.metrics import log_metric

device_arn = Devices.Amazon.SV1


@hybrid_job(device=device_arn)  # Choose priority device
def run_hybrid_job(num_tasks=1):
    device = AwsDevice(device_arn)  # Declare AwsDevice within the hybrid job

    # Create a parametric circuit
    circ = Circuit()
    circ.rx(0, FreeParameter("theta"))
    circ.cnot(0, 1)
    circ.expectation(observable=Observable.X(), target=0)

    theta = 0.0  # Initial parameter

    for i in range(num_tasks):
        task = device.run(circ, shots=100, inputs={"theta": theta})  # Input parameters
        exp_val = task.result().values[0]

        theta += exp_val  # Modify the parameter (possibly gradient descent)

        log_metric(metric_name="exp_val", value=exp_val, iteration_number=i)

    return {"final_theta": theta, "final_exp_val": exp_val}

ハイブリッドジョブを作成するには、通常の Python 関数と同様に被デコレート関数を呼び出します。ただし、デコレータ関数は、被デコレート関数の結果ではなくハイブリッドジョブのハンドルを返します。ジョブの完了後に結果を取得するには、job.result() を使用します。

job = run_hybrid_job(num_tasks=1)
result = job.result() 

@hybrid_job デコレータのデバイス引数には、ハイブリッドジョブが優先アクセス権を持つデバイスを指定します。この場合はシミュレーター SV1 です。QPU の優先順位を取得するには、デコレータで指定されたものと一致するデバイス ARN を被デコレート関数内で使用する必要があります。@hybrid_job で宣言されたデバイス ARN をキャプチャするには、ヘルパー関数 get_job_device_arn() を使用するという便利な方法もあります。

ハイパーパラメータ

run_hybrid_job() 関数が取る引数として、作成される量子タスクの数を制御できる num_tasks があります。ハイブリッドジョブは、この引数をハイパーパラメータとして自動的にキャプチャします。

メトリクスとログ記録

run_hybrid_job() 関数内では、反復アルゴリズムからのメトリクスが log_metrics を使用して記録されます。メトリクスは、Braket コンソールページのハイブリッドジョブタブに自動的にプロットされます。Braket コストトラッカーを使用することで、ハイブリッドジョブ実行中の量子タスクのコストをほぼリアルタイムで追跡できます。上記の例では、結果タイプにおける最初の確率を記録するメトリクス名「probability」(確率) を使用しています。

結果を取得する

ハイブリッドジョブが完了したら、 job.result() を使用することで、ハイブリッドジョブの結果を取得できます。return ステートメントに記述したすべてのオブジェクトが、Braket によって自動的にキャプチャされます。関数が返すオブジェクトは、各要素がシリアル化可能なタプルである必要があります。例えば、次のコードは成功している例と失敗している例を示しています。

import numpy as np


# Working example
@hybrid_job(device=Devices.Amazon.SV1)
def passing():
    np_array = np.random.rand(5)
    return np_array  # Serializable

# # Failing example
# @hybrid_job(device=Devices.Amazon.SV1)
# def failing():
#     return MyObject() # Not serializable

ジョブ名

デフォルトでは、このハイブリッドジョブの名前は関数名を使用して提案されます。また、50 文字までの長さのカスタム名を指定することもできます。例えば、次のコードでは、ジョブ名は「my-job-name」です。

@hybrid_job(device=Devices.Amazon.SV1, job_name="my-job-name")
def function():
    pass

ローカルモード

ローカルジョブは、デコレータに引数 local=True を追加することで作成されます。これにより、ラップトップなどのローカルコンピューティング環境のコンテナ化された環境でハイブリッドジョブが実行されます。ローカルジョブでは、量子タスクに対して優先的なキューイングが行われることはありません。マルチノードや MPI などの高度なケースでは、ローカルジョブが、必要な Braket 環境変数にアクセスできる場合があります。次のコードは、SV1 シミュレーターをデバイスとして使用してローカルハイブリッドジョブを作成しています。

@hybrid_job(device=Devices.Amazon.SV1, local=True)
def run_hybrid_job(num_tasks=1):
    return ...

この他にも、あらゆるハイブリッドジョブオプションがサポートされています。オプションのリストについては、「braket.jobs.quantum_job_creation module」を参照してください。

追加の Python パッケージとソースコードをインストールする

希望の Python パッケージを使用できるようにランタイム環境をカスタマイズできます。パッケージ名のリストである requirements.txt ファイルか、独自のコンテナ (BYOC) のいずれかを使用できます。例えば、requirements.txt ファイルには、インストールが必要な他のパッケージを含めることができます。

qiskit 
pennylane >= 0.31
mitiq == 0.29

requirements.txt ファイルを使用してランタイム環境をカスタマイズするには、以下のコード例を参照してください。

@hybrid_job(device=Devices.Amazon.SV1, dependencies="requirements.txt")
def run_hybrid_job(num_tasks=1):
    return ...

または、次のようにパッケージ名を Python リストとして指定することもできます。

@hybrid_job(device=Devices.Amazon.SV1, dependencies=["qiskit", "pennylane>=0.31", "mitiq==0.29"])
def run_hybrid_job(num_tasks=1):
    return ...

追加のソースコードは、モジュールのリストとして指定することも、次のコード例のように単一のモジュールとして指定することもできます。

@hybrid_job(device=Devices.Amazon.SV1, include_modules=["my_module1", "my_module2"])
def run_hybrid_job(num_tasks=1):
    return ...

ハイブリッドジョブインスタンスにおいてデータをロードしたり保存したりする

入力トレーニングデータを指定する

ハイブリッドジョブを作成する場合は、Amazon Simple Storage Service (Amazon S3) バケットを指定して入力トレーニングデータセットを提供できます。また、ローカルパスを指定すれば、データが Braket によって Amazon S3 の s3://<default_bucket_name>/jobs/<job_name>/<timestamp>/data/<channel_name> に自動的にアップロードされます。ローカルパスを指定した場合、チャネル名はデフォルトで「input」になります。次のコードは、ローカルパス data/file.npy からアップロードされた numpy ファイルを示しています。

import numpy as np


@hybrid_job(device=Devices.Amazon.SV1, input_data="data/file.npy")
def run_hybrid_job(num_tasks=1):
    data = np.load("data/file.npy")
    return ...

S3 の場合は、get_input_data_dir() ヘルパー関数を使用する必要があります。

import numpy as np
from braket.jobs import get_input_data_dir

s3_path = "s3://amazon-braket-us-east-1-123456789012/job-data/file.npy"


@hybrid_job(device=None, input_data=s3_path)
def job_s3_input():
    np.load(get_input_data_dir() + "/file.npy")


@hybrid_job(device=None, input_data={"channel": s3_path})
def job_s3_input_channel():
    np.load(get_input_data_dir("channel") + "/file.npy")

チャネル値と S3 URI またはローカルパスで構成されるディクショナリを提供することで、複数の入力データソースを指定できます。

import numpy as np
from braket.jobs import get_input_data_dir

input_data = {
    "input": "data/file.npy",
    "input_2": "s3://amzn-s3-demo-bucket/data.json"
}


@hybrid_job(device=None, input_data=input_data)
def multiple_input_job():
    np.load(get_input_data_dir("input") + "/file.npy")
    np.load(get_input_data_dir("input_2") + "/data.json")

結果を S3 に保存する

被デコレート関数の return ステートメントに指定されていない結果を保存するには、すべてのファイル書き込みオペレーションに正しいディレクトリを追加する必要があります。次の例は、numpy 配列と matplotlib 図の保存を示しています。

import matplotlib.pyplot as plt
import numpy as np


@hybrid_job(device=Devices.Amazon.SV1)
def run_hybrid_job(num_tasks=1):
    result = np.random.rand(5)

    # Save a numpy array
    np.save("result.npy", result)

    # Save a matplotlib figure
    plt.plot(result)
    plt.savefig("fig.png")
    return ...

すべての結果は model.tar.gz という名前のファイルに圧縮されます。この結果をダウンロードするには、Python 関数 job.result() を使用するか、または Braket マネジメントコンソールのハイブリッドジョブページにある結果フォルダに移動します。

チェックポイントからの保存と再開

長時間実行されるハイブリッドジョブの場合、アルゴリズムの中間状態を定期的に保存することをお勧めします。組み込みの save_job_checkpoint() ヘルパー関数を使用するか、ファイルを AMZN_BRAKET_JOB_RESULTS_DIR パスに保存できます。後者のパスはヘルパー関数 get_job_results_dir() で使用できます。

以下は、最小限の作業によって、ハイブリッドジョブデコレータを使用してチェックポイントを保存およびロードする例です。

from braket.jobs import save_job_checkpoint, load_job_checkpoint, hybrid_job


@hybrid_job(device=None, wait_until_complete=True)
def function():
    save_job_checkpoint({"a": 1})


job = function()
job_name = job.name
job_arn = job.arn


@hybrid_job(device=None, wait_until_complete=True, copy_checkpoints_from_job=job_arn)
def continued_function():
    load_job_checkpoint(job_name)


continued_job = continued_function()

最初のハイブリッドジョブでは、 save_job_checkpoint() が、保存するデータを含むディクショナリを指定して呼び出されています。デフォルトでは、すべての値がテキストとしてシリアル化されます。numpy 配列など、より複雑な Python オブジェクトのチェックポイントファイルを作成するには、data_format = PersistedJobDataFormat.PICKLED_V4 を設定します。このコードは、「checkpoints」というサブフォルダにあるハイブリッドジョブアーティファクトに、デフォルト名 <jobname>.json を持つチェックポイントファイルを作成および上書きします。

チェックポイントから再開する新しいハイブリッドジョブを作成するには、copy_checkpoints_from_job=job_arn を渡す必要があります。ここで、job_arn は直前のジョブのハイブリッドジョブ ARN です。次に、load_job_checkpoint(job_name) を使用してチェックポイントの内容をロードします。

ハイブリッドジョブデコレータに関するベストプラクティス

非同期性を採用する

デコレータ注釈を使用して作成されたハイブリッドジョブは、古典的リソースと量子リソースが利用可能になると実行され、非同期です。アルゴリズムの進行状況をモニタリングするには、Braket Management Console または Amazon CloudWatch を使用します。実行すべきアルゴリズムを送信すると、Braket により、スケーラブルなコンテナ化された環境で実行され、完了すると結果が取得されます。

反復変分アルゴリズムを実行する

ハイブリッドジョブでは、反復量子古典アルゴリズムを実行するツールを使用できます。純粋に量子力学的な問題には、量子タスクまたは量子タスクのバッチを使用します。特定の QPU への優先的なアクセスは、古典的な処理で QPU への複数の反復呼び出しを断続的に必要とする、長時間実行される変分アルゴリズムに最も有益です。

ローカルモードを使用してデバッグする

QPU でハイブリッドジョブを実行する場合は、まずシミュレーター SV1 で実行して、期待どおりに実行されるかを確認することをお勧めします。小規模なテストでは、ローカルモードで実行することで迅速な反復実行とデバッグを行うことができます。

独自のコンテナ (BYOC) により再現性を向上する

コンテナ化された環境内にソフトウェアとその依存関係をカプセル化することで、再現可能な実験を作成します。すべてのコード、依存関係、設定をコンテナにパッケージ化することで、潜在的な競合やバージョニングの問題を回避できます。

マルチインスタンス分散シミュレーター

多数の回路を実行するには、組み込みの MPI サポートを使用して、単一のハイブリッドジョブにおいて、複数のインスタンスでローカルシミュレーターを実行することを検討してください。詳細については、「embedded simulators」を参照してください。

パラメトリック回路を使用する

ハイブリッドジョブから送信したパラメトリック回路は、パラメトリックコンパイルを使用して特定の QPU で自動的にコンパイルされるため、アルゴリズムの実行時間が短縮されます。

定期的にチェックポイントを作成する

長時間実行されるハイブリッドジョブの場合、アルゴリズムの中間状態を定期的に保存することをお勧めします。

その他の例、ユースケース、ベストプラクティスについては、GitHub の「Amazon Braket examples」を参照してください。