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)