braket.jobs.hybrid_job module

braket.jobs.hybrid_job.hybrid_job(*, device: str | None, include_modules: str | ModuleType | Iterable[str | ModuleType] | None = None, dependencies: str | Path | list[str] | None = None, local: bool = False, job_name: str | None = None, image_uri: str | None = None, input_data: str | dict | S3DataSourceConfig | None = None, wait_until_complete: bool = False, instance_config: InstanceConfig | None = None, distribution: str | None = None, copy_checkpoints_from_job: str | None = None, checkpoint_config: CheckpointConfig | None = None, role_arn: str | None = None, stopping_condition: StoppingCondition | None = None, output_data_config: OutputDataConfig | None = None, aws_session: AwsSession | None = None, tags: dict[str, str] | None = None, logger: Logger = <Logger braket.jobs.hybrid_job (WARNING)>, quiet: bool | None = None, reservation_arn: str | None = None) → Callable[source]

Defines a hybrid job by decorating the entry point function. The job will be created when the decorated function is called.

The job created will be a LocalQuantumJob when local is set to True, otherwise an AwsQuantumJob. The following parameters will be ignored when running a job with local set to True: wait_until_complete, instance_config, distribution, copy_checkpoints_from_job, stopping_condition, tags, logger, and quiet.

Parameters:

  • device (str | None) – Device ARN of the QPU device that receives priority quantum task queueing once the hybrid job begins running. Each QPU has a separate hybrid jobs queue so that only one hybrid job is running at a time. The device string is accessible in the hybrid job instance as the environment variable “AMZN_BRAKET_DEVICE_ARN”. When using embedded simulators, you may provide the device argument as string of the form: “local:<provider>/<simulator_name>” or None.

  • include_modules (str | ModuleType | Iterable[str | ModuleType] | None) – Either a single module or module name or a list of module or module names referring to local modules to be included. Any references to members of these modules in the hybrid job algorithm code will be serialized as part of the algorithm code. Default: []

  • dependencies (str | Path | list[str] | None) – Path (absolute or relative) to a requirements.txt file, or alternatively a list of strings, with each string being a requirement specifier, to be used for the hybrid job.

  • local (bool) – Whether to use local mode for the hybrid job. Default: False

  • job_name (str | None) – A string that specifies the name with which the job is created. Allowed pattern for job name: ^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,50}$. Defaults to f’{decorated-function-name}-{timestamp}’.

  • image_uri (str | None) – A str that specifies the ECR image to use for executing the job. retrieve_image() function may be used for retrieving the ECR image URIs for the containers supported by Braket. Default: <Braket base image_uri>.

  • input_data (str | dict | S3DataSourceConfig | None) – Information about the training data. Dictionary maps channel names to local paths or S3 URIs. Contents found at any local paths will be uploaded to S3 at f’s3://{default_bucket_name}/jobs/{job_name}/data/{channel_name}’. If a local path, S3 URI, or S3DataSourceConfig is provided, it will be given a default channel name “input”. Default: {}.

  • wait_until_complete (bool) – True if we should wait until the job completes. This would tail the job logs as it waits. Otherwise False. Ignored if using local mode. Default: False.

  • instance_config (InstanceConfig | None) – Configuration of the instance(s) for running the classical code for the hybrid job. Default: InstanceConfig(instanceType='ml.m5.large', instanceCount=1, volumeSizeInGB=30).

  • distribution (str | None) – A str that specifies how the job should be distributed. If set to “data_parallel”, the hyperparameters for the job will be set to use data parallelism features for PyTorch or TensorFlow. Default: None.

  • copy_checkpoints_from_job (str | None) – A str that specifies the job ARN whose checkpoint you want to use in the current job. Specifying this value will copy over the checkpoint data from use_checkpoints_from_job’s checkpoint_config s3Uri to the current job’s checkpoint_config s3Uri, making it available at checkpoint_config.localPath during the job execution. Default: None

  • checkpoint_config (CheckpointConfig | None) – Configuration that specifies the location where checkpoint data is stored. Default: CheckpointConfig(localPath='/opt/jobs/checkpoints', s3Uri=f's3://{default_bucket_name}/jobs/{job_name}/checkpoints').

  • role_arn (str | None) – A str providing the IAM role ARN used to execute the script. Default: IAM role returned by AwsSession’s get_default_jobs_role().

  • stopping_condition (StoppingCondition | None) – The maximum length of time, in seconds, and the maximum number of tasks that a job can run before being forcefully stopped. Default: StoppingCondition(maxRuntimeInSeconds=5 * 24 * 60 * 60).

  • output_data_config (OutputDataConfig | None) – Specifies the location for the output of the job. Default: OutputDataConfig(s3Path=f's3://{default_bucket_name}/jobs/{job_name}/data', kmsKeyId=None).

  • aws_session (AwsSession | None) – AwsSession for connecting to AWS Services. Default: AwsSession()

  • tags (dict[str, str] | None) – Dict specifying the key-value pairs for tagging this job. Default: {}.

  • logger (Logger) – Logger object with which to write logs, such as task statuses while waiting for task to be in a terminal state. Default: getLogger(__name__)

  • quiet (bool | None) – Sets the verbosity of the logger to low and does not report queue position. Default is False.

  • reservation_arn (str | None) – the reservation window arn provided by Braket Direct to reserve exclusive usage for the device to run the hybrid job on. Default: None.

Returns:

Callable – the callable for creating a Hybrid Job.