def test_attach_custom_image(sagemaker_session):
training_image = "1.dkr.ecr.us-west-2.amazonaws.com/my_custom_xgboost_image:latest"
returned_job_description = {
"AlgorithmSpecification": {"TrainingInputMode": "File", "TrainingImage": training_image},
"HyperParameters": {
"sagemaker_submit_directory": '"s3://some/sourcedir.tar.gz"',
"sagemaker_program": '"iris-dnn-classifier.py"',
"sagemaker_s3_uri_training": '"sagemaker-3/integ-test-data/tf_iris"',
"sagemaker_container_log_level": '"logging.INFO"',
"sagemaker_job_name": '"neo"',
"training_steps": "100",
"sagemaker_region": '"us-west-2"',
},
"RoleArn": "arn:aws:iam::366:role/SageMakerRole",
"ResourceConfig": {
"VolumeSizeInGB": 30,
"InstanceCount": 1,
"InstanceType": "ml.c4.xlarge",
},
"StoppingCondition": {"MaxRuntimeInSeconds": 24 * 60 * 60},
"TrainingJobName": "neo",
"TrainingJobStatus": "Completed",
"TrainingJobArn": "arn:aws:sagemaker:us-west-2:336:training-job/neo",
"OutputDataConfig": {"KmsKeyId": "", "S3OutputPath": "s3://place/output/neo"},
"TrainingJobOutput": {"S3TrainingJobOutput": "s3://here/output.tar.gz"},
}
sagemaker_session.sagemaker_client.describe_training_job = Mock(
name="describe_training_job", return_value=returned_job_description
)
with pytest.raises(TypeError) as error:
XGBoost.attach(training_job_name="neo", sagemaker_session=sagemaker_session)
assert "expected string" in str(error)
def test_attach_wrong_framework(sagemaker_session):
rjd = {
"AlgorithmSpecification": {
"TrainingInputMode": "File",
"TrainingImage": "1.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-py3-cpu:1.0.4",
},
"HyperParameters": {
"sagemaker_submit_directory": '"s3://some/sourcedir.tar.gz"',
"checkpoint_path": '"s3://other/1508872349"',
"sagemaker_program": '"iris-dnn-classifier.py"',
"sagemaker_container_log_level": '"logging.INFO"',
"training_steps": "100",
"sagemaker_region": '"us-west-2"',
},
"RoleArn": "arn:aws:iam::366:role/SageMakerRole",
"ResourceConfig": {
"VolumeSizeInGB": 30,
"InstanceCount": 1,
"InstanceType": "ml.c4.xlarge",
},
"StoppingCondition": {"MaxRuntimeInSeconds": 24 * 60 * 60},
"TrainingJobName": "neo",
"TrainingJobStatus": "Completed",
"TrainingJobArn": "arn:aws:sagemaker:us-west-2:336:training-job/neo",
"OutputDataConfig": {"KmsKeyId": "", "S3OutputPath": "s3://place/output/neo"},
"TrainingJobOutput": {"S3TrainingJobOutput": "s3://here/output.tar.gz"},
}
sagemaker_session.sagemaker_client.describe_training_job = Mock(
name="describe_training_job", return_value=rjd
)
with pytest.raises(ValueError) as error:
XGBoost.attach(training_job_name="neo", sagemaker_session=sagemaker_session)
assert "didn't use image for requested framework" in str(error)
def test_attach(sagemaker_session, xgboost_version):
training_image = "1.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:{}-cpu-{}".format(
xgboost_version, PYTHON_VERSION)
returned_job_description = {
"AlgorithmSpecification": {
"TrainingInputMode": "File",
"TrainingImage": training_image
},
"HyperParameters": {
"sagemaker_submit_directory": '"s3://some/sourcedir.tar.gz"',
"sagemaker_program": '"iris-dnn-classifier.py"',
"sagemaker_s3_uri_training":
'"sagemaker-3/integ-test-data/tf_iris"',
"sagemaker_enable_cloudwatch_metrics": "false",
"sagemaker_container_log_level": '"logging.INFO"',
"sagemaker_job_name": '"neo"',
"training_steps": "100",
"sagemaker_region": '"us-west-2"',
},
"RoleArn": "arn:aws:iam::366:role/SageMakerRole",
"ResourceConfig": {
"VolumeSizeInGB": 30,
"InstanceCount": 1,
"InstanceType": "ml.c4.xlarge",
},
"StoppingCondition": {
"MaxRuntimeInSeconds": 24 * 60 * 60
},
"TrainingJobName": "neo",
"TrainingJobStatus": "Completed",
"TrainingJobArn": "arn:aws:sagemaker:us-west-2:336:training-job/neo",
"OutputDataConfig": {
"KmsKeyId": "",
"S3OutputPath": "s3://place/output/neo"
},
"TrainingJobOutput": {
"S3TrainingJobOutput": "s3://here/output.tar.gz"
},
}
sagemaker_session.sagemaker_client.describe_training_job = Mock(
name="describe_training_job", return_value=returned_job_description)
estimator = XGBoost.attach(training_job_name="neo",
sagemaker_session=sagemaker_session)
assert estimator._current_job_name == "neo"
assert estimator.latest_training_job.job_name == "neo"
assert estimator.py_version == PYTHON_VERSION
assert estimator.framework_version == xgboost_version
assert estimator.role == "arn:aws:iam::366:role/SageMakerRole"
assert estimator.train_instance_count == 1
assert estimator.train_max_run == 24 * 60 * 60
assert estimator.input_mode == "File"
assert estimator.base_job_name == "neo"
assert estimator.output_path == "s3://place/output/neo"
assert estimator.output_kms_key == ""
assert estimator.hyperparameters()["training_steps"] == "100"
assert estimator.source_dir == "s3://some/sourcedir.tar.gz"
assert estimator.entry_point == "iris-dnn-classifier.py"
