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)
