def test_distributed_training(strftime, sagemaker_session,
xgboost_framework_version):
xgboost = XGBoost(
entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
instance_count=DIST_INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
py_version=PYTHON_VERSION,
framework_version=xgboost_framework_version,
)
inputs = "s3://mybucket/train"
xgboost.fit(inputs=inputs)
sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
assert sagemaker_call_names == ["train", "logs_for_job"]
boto_call_names = [
c[0] for c in sagemaker_session.boto_session.method_calls
]
assert boto_call_names == ["resource"]
expected_train_args = _create_train_job(xgboost_framework_version,
DIST_INSTANCE_COUNT)
expected_train_args["input_config"][0]["DataSource"]["S3DataSource"][
"S3Uri"] = inputs
actual_train_args = sagemaker_session.method_calls[0][2]
assert actual_train_args == expected_train_args
model = xgboost.create_model()
expected_image_base = "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:{}-cpu-{}"
assert {
"Environment": {
"SAGEMAKER_SUBMIT_DIRECTORY":
"s3://mybucket/sagemaker-xgboost-{}/source/sourcedir.tar.gz".
format(TIMESTAMP),
"SAGEMAKER_PROGRAM":
"dummy_script.py",
"SAGEMAKER_REGION":
"us-west-2",
"SAGEMAKER_CONTAINER_LOG_LEVEL":
"20",
},
"Image":
expected_image_base.format(xgboost_framework_version, PYTHON_VERSION),
"ModelDataUrl":
"s3://m/m.tar.gz",
} == model.prepare_container_def(CPU)
assert "cpu" in model.prepare_container_def(CPU)["Image"]
predictor = xgboost.deploy(1, CPU)
assert isinstance(predictor, XGBoostPredictor)
def test_xgboost_gpu(time, strftime, sagemaker_session,
xgboost_gpu_framework_version):
xgboost = XGBoost(
entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
instance_type=GPU_INSTANCE_TYPE,
instance_count=1,
framework_version=xgboost_gpu_framework_version,
)
inputs = "s3://mybucket/train"
xgboost.fit(inputs=inputs, experiment_config=EXPERIMENT_CONFIG)
sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
assert sagemaker_call_names == ["train", "logs_for_job"]
boto_call_names = [
c[0] for c in sagemaker_session.boto_session.method_calls
]
assert boto_call_names == ["resource"]
expected_train_args = _create_train_job(xgboost_gpu_framework_version,
instance_type=GPU_INSTANCE_TYPE)
expected_train_args["input_config"][0]["DataSource"]["S3DataSource"][
"S3Uri"] = inputs
expected_train_args["experiment_config"] = EXPERIMENT_CONFIG
actual_train_args = sagemaker_session.method_calls[0][2]
assert actual_train_args == expected_train_args
model = xgboost.create_model()
assert {
"Environment": {
"SAGEMAKER_SUBMIT_DIRECTORY":
"s3://mybucket/sagemaker-xgboost-{}/source/sourcedir.tar.gz".
format(TIMESTAMP),
"SAGEMAKER_PROGRAM":
"dummy_script.py",
"SAGEMAKER_REGION":
"us-west-2",
"SAGEMAKER_CONTAINER_LOG_LEVEL":
"20",
},
"Image": _get_full_image_uri(xgboost_gpu_framework_version),
"ModelDataUrl": "s3://m/m.tar.gz",
} == model.prepare_container_def(GPU_INSTANCE_TYPE)
predictor = xgboost.deploy(1, GPU_INSTANCE_TYPE)
assert isinstance(predictor, XGBoostPredictor)
