def test_copy_to(mock_run, mock_check_output, repo_uri, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" source_path = str(Path("test", "source", "dir", "path", "srcfile.txt")) dest_path = str(Path("test", "dest", "dir", "path", "dstfile.txt")) mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), str.encode(""), str.encode(""), ] with _LocalJobContainer(image_uri, aws_session) as container: container.copy_to(source_path, dest_path) mock_check_output.assert_any_call(["docker", "images", "-q", image_uri]) mock_check_output.assert_any_call([ "docker", "run", "-d", "--rm", local_image_name, "tail", "-f", "/dev/null" ]) mock_check_output.assert_any_call([ "docker", "exec", running_container_name, "mkdir", "-p", str(Path("test", "dest", "dir", "path")), ]) mock_check_output.assert_any_call( ["docker", "cp", source_path, f"{running_container_name}:{dest_path}"]) assert mock_check_output.call_count == 4 mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 1
def test_pull_container_forced_update_invalid_name(mock_run, mock_check_output, repo_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" mock_logger = Mock() mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), ] with _LocalJobContainer(image_uri=local_image_name, aws_session=aws_session, force_update=True, logger=mock_logger): pass mock_logger.warning.assert_called_with( f"Unable to update {local_image_name}.") mock_check_output.assert_any_call( ["docker", "images", "-q", local_image_name]) mock_check_output.assert_any_call([ "docker", "run", "-d", "--rm", local_image_name, "tail", "-f", "/dev/null" ]) assert mock_check_output.call_count == 2 mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 1
def test_pull_container( mock_run, mock_check_output, repo_uri, image_uri, aws_session, forced_update, check_output, local_image_name, ): running_container_name = "RunningContainer" test_token = "Test Token" mock_check_output.side_effect = check_output aws_session.ecr_client.get_authorization_token.return_value = { "authorizationData": [{ "authorizationToken": base64.b64encode(str.encode(test_token)) }] } with _LocalJobContainer(image_uri=image_uri, aws_session=aws_session, force_update=forced_update): pass mock_check_output.assert_any_call(["docker", "images", "-q", image_uri]) mock_check_output.assert_any_call([ "docker", "run", "-d", "--rm", local_image_name, "tail", "-f", "/dev/null" ]) assert mock_check_output.call_count == len(check_output) mock_run.assert_any_call( ["docker", "login", "-u", "AWS", "-p", test_token, repo_uri]) mock_run.assert_any_call(["docker", "pull", image_uri]) mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 3
def test_pull_fails_invalid_uri(mock_run, mock_check_output, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" mock_check_output.side_effect = [ str.encode(""), str.encode(local_image_name), str.encode(running_container_name), ] aws_session.ecr_client.get_authorization_token.return_value = {} with _LocalJobContainer("TestURI", aws_session): pass
def test_make_dir_fails(mock_run, mock_check_output, repo_uri, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" test_dir_path = "/test/dir/path" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), subprocess.CalledProcessError("Test Error", "test", str.encode("test output")), ] with _LocalJobContainer(image_uri, aws_session) as container: container.makedir(test_dir_path)
def test_copy_from_fails(mock_run, mock_check_output, repo_uri, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" source_path = "/test/source/dir/path/srcfile.txt" dest_path = "/test/dest/dir/path/dstfile.txt" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), subprocess.CalledProcessError("Test Error", "test", str.encode("test output")), ] with _LocalJobContainer(image_uri, aws_session) as container: container.copy_from(source_path, dest_path)
def test_customer_script_check_output(mock_run, mock_check_output, repo_uri, image_uri, aws_session): mock_logger = Mock() local_image_name = "LocalImageName" running_container_name = "RunningContainer" env_variables = { "ENV0": "VALUE0", "ENV1": "VALUE1", } run_program_name = "Run Program Name" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), str.encode(run_program_name), ] run_result = CompletedProcess(None, 0, str.encode("Test Output"), str.encode("Test Error")) mock_run.side_effect = [run_result, None] with _LocalJobContainer(image_uri, aws_session, mock_logger) as container: container.run_local_job(env_variables) assert container.run_result == run_result mock_check_output.assert_any_call(["docker", "images", "-q", image_uri]) mock_check_output.assert_any_call([ "docker", "run", "-d", "--rm", local_image_name, "tail", "-f", "/dev/null" ]) mock_check_output.assert_any_call([ "docker", "exec", running_container_name, "printenv", "SAGEMAKER_PROGRAM" ]) assert mock_check_output.call_count == 3 mock_run.assert_any_call( [ "docker", "exec", "-w", "/opt/ml/code/", "-e", "ENV0=VALUE0", "-e", "ENV1=VALUE1", running_container_name, "python", run_program_name, ], capture_output=True, ) mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 2 mock_logger.error.assert_called_with("Test Error")
def test_pull_fails_unknown_reason(mock_run, mock_check_output, repo_uri, image_uri, aws_session): test_token = "Test Token" mock_check_output.side_effect = [ str.encode(""), str.encode(""), ] aws_session.ecr_client.get_authorization_token.return_value = { "authorizationData": [{ "authorizationToken": base64.b64encode(str.encode(test_token)) }] } with _LocalJobContainer(image_uri, aws_session): pass
def test_run_fails_no_program(mock_run, mock_check_output, repo_uri, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" env_variables = { "ENV0": "VALUE0", "ENV1": "VALUE1", } mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), str.encode(""), ] with _LocalJobContainer(image_uri, aws_session) as container: container.run_local_job(env_variables)
def test_start_and_stop(mock_run, mock_check_output, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), ] with _LocalJobContainer(image_uri, aws_session): pass mock_check_output.assert_any_call(["docker", "images", "-q", image_uri]) mock_check_output.assert_any_call([ "docker", "run", "-d", "--rm", local_image_name, "tail", "-f", "/dev/null" ]) assert mock_check_output.call_count == 2 mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 1
def test_run_job_success(mock_run, mock_check_output, repo_uri, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" env_variables = { "ENV0": "VALUE0", "ENV1": "VALUE1", } run_program_name = "Run Program Name" expected_run_output = "Expected Run Output" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), str.encode(run_program_name), str.encode(expected_run_output), ] with _LocalJobContainer(image_uri, aws_session) as container: container.run_local_job(env_variables) run_output = container.run_log assert run_output == expected_run_output mock_check_output.assert_any_call(["docker", "images", "-q", image_uri]) mock_check_output.assert_any_call([ "docker", "run", "-d", "--rm", local_image_name, "tail", "-f", "/dev/null" ]) mock_check_output.assert_any_call([ "docker", "exec", running_container_name, "printenv", "SAGEMAKER_PROGRAM" ]) mock_check_output.assert_any_call([ "docker", "exec", "-w", "/opt/ml/code/", "-e", "ENV0=VALUE0", "-e", "ENV1=VALUE1", running_container_name, "python", run_program_name, ]) assert mock_check_output.call_count == 4 mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 1
def test_make_dir(mock_run, mock_check_output, repo_uri, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" test_dir_path = "/test/dir/path" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), str.encode(""), ] with _LocalJobContainer(image_uri, aws_session) as container: container.makedir(test_dir_path) mock_check_output.assert_any_call(["docker", "images", "-q", image_uri]) mock_check_output.assert_any_call([ "docker", "run", "-d", "--rm", local_image_name, "tail", "-f", "/dev/null" ]) mock_check_output.assert_any_call([ "docker", "exec", running_container_name, "mkdir", "-p", test_dir_path ]) assert mock_check_output.call_count == 3 mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 1
def test_copy_from(mock_run, mock_check_output, repo_uri, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" source_path = "/test/source/dir/path/srcfile.txt" dest_path = "/test/dest/dir/path/dstfile.txt" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), str.encode(""), str.encode(""), ] with _LocalJobContainer(image_uri, aws_session) as container: container.copy_from(source_path, dest_path) mock_check_output.assert_any_call(["docker", "images", "-q", image_uri]) mock_check_output.assert_any_call([ "docker", "run", "-d", "--rm", local_image_name, "tail", "-f", "/dev/null" ]) mock_check_output.assert_any_call( ["docker", "cp", f"{running_container_name}:{source_path}", dest_path]) assert mock_check_output.call_count == 3 mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 1
def test_run_customer_script_fails(mock_run, mock_check_output, repo_uri, image_uri, aws_session): mock_logger = Mock() local_image_name = "LocalImageName" running_container_name = "RunningContainer" env_variables = { "ENV0": "VALUE0", "ENV1": "VALUE1", } run_program_name = "Run Program Name" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), str.encode(run_program_name), ] expected_exception = Exception("Test Error") mock_run.side_effect = [expected_exception, None] with _LocalJobContainer(image_uri, aws_session, mock_logger) as container: container.run_local_job(env_variables) assert container.run_result == expected_exception assert mock_check_output.call_count == 3 mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 2 mock_logger.error.assert_called_with(expected_exception)
def test_customer_script_fails(mock_run, mock_check_output, repo_uri, image_uri, aws_session): local_image_name = "LocalImageName" running_container_name = "RunningContainer" env_variables = { "ENV0": "VALUE0", "ENV1": "VALUE1", } run_program_name = "Run Program Name" expected_error_output = "Expected Error Output" mock_check_output.side_effect = [ str.encode(local_image_name), str.encode(running_container_name), str.encode(run_program_name), subprocess.CalledProcessError("Test Error", "test", str.encode(expected_error_output)), ] with _LocalJobContainer(image_uri, aws_session) as container: container.run_local_job(env_variables) run_output = container.run_log assert run_output == expected_error_output assert mock_check_output.call_count == 4 mock_run.assert_any_call(["docker", "stop", running_container_name]) assert mock_run.call_count == 1
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)