def test_no_data_input(container, aws_session, creation_kwargs, input_data_config): input_data_config.append({ # this channel won't match any data "channelName": "no-data", "dataSource": { "s3DataSource": { "s3Uri": "s3://input_bucket/irrelevant" } }, }) creation_kwargs.update({"inputDataConfig": input_data_config}) no_data_found = "No data found for channel 'no-data'" with pytest.raises(RuntimeError, match=no_data_found): setup_container(container, aws_session, **creation_kwargs)
def test_duplicate_input(container, aws_session, creation_kwargs, input_data_config): input_data_config.append({ # this is a duplicate channel "channelName": "single-file", "dataSource": { "s3DataSource": { "s3Uri": "s3://input_bucket/irrelevant" } }, }) creation_kwargs.update({"inputDataConfig": input_data_config}) dupes_not_allowed = "Duplicate channel names not allowed for input data: single-file" with pytest.raises(ValueError, match=dupes_not_allowed): setup_container(container, aws_session, **creation_kwargs)
def test_basic_setup(container, aws_session, creation_kwargs, expected_envs): aws_session.parse_s3_uri.return_value = ["test_bucket", "test_location"] envs = setup_container(container, aws_session, **creation_kwargs) assert envs == expected_envs container.makedir.assert_any_call("/opt/ml/model") container.makedir.assert_any_call( expected_envs["AMZN_BRAKET_CHECKPOINT_DIR"]) assert container.makedir.call_count == 2
def test_temporary_credentials(container, aws_session, creation_kwargs, expected_envs): aws_session.boto_session.get_credentials.return_value.token = "Test Token" expected_envs["AWS_SESSION_TOKEN"] = "Test Token" aws_session.parse_s3_uri.return_value = ["test_bucket", "test_location"] envs = setup_container(container, aws_session, **creation_kwargs) assert envs == expected_envs container.makedir.assert_any_call("/opt/ml/model") container.makedir.assert_any_call( expected_envs["AMZN_BRAKET_CHECKPOINT_DIR"]) assert container.makedir.call_count == 2
def test_input(container, aws_session, creation_kwargs, input_data_config): creation_kwargs.update({"inputDataConfig": input_data_config}) setup_container(container, aws_session, **creation_kwargs) download_locations = [ call[0][1] for call in aws_session.download_from_s3.call_args_list ] expected_downloads = [ Path("single-file", "file-1.txt"), Path("directory-no-slash", "file-1.txt"), Path("directory-no-slash", "file-2.txt"), Path("directory-slash", "file-1.txt"), Path("directory-slash", "file-2.txt"), Path("directory-prefix", "input-dir", "file-1.txt"), Path("directory-prefix", "input-dir", "file-2.txt"), Path("files-prefix", "file-1.txt"), Path("files-prefix", "file-2.txt"), ] for download, expected_download in zip(download_locations, expected_downloads): assert download.endswith(str(expected_download))
def test_compressed_script_mode(container, aws_session, creation_kwargs, expected_envs, compressed_script_mode_config): creation_kwargs["algorithmSpecification"] = compressed_script_mode_config expected_envs[ "AMZN_BRAKET_SCRIPT_S3_URI"] = "s3://amazon-braket-jobs/job-path/my_archive.gzip" expected_envs["AMZN_BRAKET_SCRIPT_COMPRESSION_TYPE"] = "gzip" aws_session.parse_s3_uri.return_value = ["test_bucket", "test_location"] envs = setup_container(container, aws_session, **creation_kwargs) assert envs == expected_envs container.makedir.assert_any_call("/opt/ml/model") container.makedir.assert_any_call( expected_envs["AMZN_BRAKET_CHECKPOINT_DIR"]) assert container.makedir.call_count == 2
def test_hyperparameters(tempfile, json, container, aws_session, creation_kwargs, expected_envs): with patch("builtins.open", mock_open()): tempfile.return_value.__enter__.return_value = "temporaryDir" creation_kwargs["hyperParameters"] = {"test": "hyper"} expected_envs[ "AMZN_BRAKET_HP_FILE"] = "/opt/braket/input/config/hyperparameters.json" aws_session.parse_s3_uri.return_value = [ "test_bucket", "test_location" ] envs = setup_container(container, aws_session, **creation_kwargs) assert envs == expected_envs container.makedir.assert_any_call("/opt/ml/model") container.makedir.assert_any_call( expected_envs["AMZN_BRAKET_CHECKPOINT_DIR"]) assert container.makedir.call_count == 2 container.copy_to.assert_called_with( os.path.join("temporaryDir", "hyperparameters.json"), "/opt/ml/input/config/hyperparameters.json", )
def create( cls, device: str, source_module: str, entry_point: str = None, image_uri: str = None, job_name: str = None, code_location: str = None, role_arn: str = None, hyperparameters: Dict[str, Any] = None, input_data: Union[str, Dict, S3DataSourceConfig] = None, output_data_config: OutputDataConfig = None, checkpoint_config: CheckpointConfig = None, aws_session: AwsSession = None, ) -> LocalQuantumJob: """Creates and runs job by setting up and running the customer script in a local docker container. Args: device (str): ARN for the AWS device which is primarily accessed for the execution of this job. source_module (str): Path (absolute, relative or an S3 URI) to a python module to be tarred and uploaded. If `source_module` is an S3 URI, it must point to a tar.gz file. Otherwise, source_module may be a file or directory. entry_point (str): A str that specifies the entry point of the job, relative to the source module. The entry point must be in the format `importable.module` or `importable.module:callable`. For example, `source_module.submodule:start_here` indicates the `start_here` function contained in `source_module.submodule`. If source_module is an S3 URI, entry point must be given. Default: source_module's name image_uri (str): A str that specifies the ECR image to use for executing the job. `image_uris.retrieve_image()` function may be used for retrieving the ECR image URIs for the containers supported by Braket. Default = `<Braket base image_uri>`. job_name (str): A str that specifies the name with which the job is created. Default: f'{image_uri_type}-{timestamp}'. code_location (str): The S3 prefix URI where custom code will be uploaded. Default: f's3://{default_bucket_name}/jobs/{job_name}/script'. role_arn (str): This field is currently not used for local jobs. Local jobs will use the current role's credentials. This may be subject to change. hyperparameters (Dict[str, Any]): Hyperparameters accessible to the job. The hyperparameters are made accessible as a Dict[str, str] to the job. For convenience, this accepts other types for keys and values, but `str()` is called to convert them before being passed on. Default: None. input_data (Union[str, S3DataSourceConfig, dict]): 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: {}. output_data_config (OutputDataConfig): Specifies the location for the output of the job. Default: OutputDataConfig(s3Path=f's3://{default_bucket_name}/jobs/{job_name}/data', kmsKeyId=None). checkpoint_config (CheckpointConfig): 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'). aws_session (AwsSession): AwsSession for connecting to AWS Services. Default: AwsSession() Returns: LocalQuantumJob: The representation of a local Braket Job. """ create_job_kwargs = prepare_quantum_job( device=device, source_module=source_module, entry_point=entry_point, image_uri=image_uri, job_name=job_name, code_location=code_location, role_arn=role_arn, hyperparameters=hyperparameters, input_data=input_data, output_data_config=output_data_config, checkpoint_config=checkpoint_config, aws_session=aws_session, ) job_name = create_job_kwargs["jobName"] if os.path.isdir(job_name): raise ValueError( f"A local directory called {job_name} already exists. " f"Please use a different job name." ) session = aws_session or AwsSession() algorithm_specification = create_job_kwargs["algorithmSpecification"] if "containerImage" in algorithm_specification: image_uri = algorithm_specification["containerImage"]["uri"] else: image_uri = retrieve_image(Framework.BASE, session.region) with _LocalJobContainer(image_uri) as container: env_variables = setup_container(container, session, **create_job_kwargs) container.run_local_job(env_variables) container.copy_from("/opt/ml/model", job_name) with open(os.path.join(job_name, "log.txt"), "w") as log_file: log_file.write(container.run_log) if "checkpointConfig" in create_job_kwargs: checkpoint_config = create_job_kwargs["checkpointConfig"] if "localPath" in checkpoint_config: checkpoint_path = checkpoint_config["localPath"] container.copy_from(checkpoint_path, os.path.join(job_name, "checkpoints")) run_log = container.run_log return LocalQuantumJob(f"local:job/{job_name}", run_log)