How to use the sagemaker.mxnet.MXNet function in sagemaker

def test_deploy(sagemaker_session, tf_version):
    estimator = MXNet(
        entry_point=SCRIPT,
        source_dir=SOURCE_DIR,
        role=ROLE,
        framework_version=tf_version,
        train_instance_count=2,
        train_instance_type=INSTANCE_TYPE_GPU,
        sagemaker_session=sagemaker_session,
        base_job_name="test-cifar",
    )

    estimator.fit("s3://mybucket/train")
    print("job succeeded: {}".format(estimator.latest_training_job.name))

    estimator.deploy(initial_instance_count=1, instance_type=INSTANCE_TYPE_CPU)
    image = IMAGE_URI_FORMAT_STRING.format(REGION, CPU_IMAGE_NAME, tf_version, "cpu", "py2")
    sagemaker_session.create_model.assert_called_with(

def test_mxnet(strftime, sagemaker_session, mxnet_version, skip_if_mms_version):
    mx = MXNet(
        entry_point=SCRIPT_PATH,
        role=ROLE,
        sagemaker_session=sagemaker_session,
        train_instance_count=INSTANCE_COUNT,
        train_instance_type=INSTANCE_TYPE,
        framework_version=mxnet_version,
    )

    inputs = "s3://mybucket/train"

    mx.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"]

def test_linear_regression(docker_image, sagemaker_local_session, local_instance_type,
                           framework_version, tmpdir):
    lr_path = os.path.join(RESOURCE_PATH, 'linear_regression')

    mx = MXNet(entry_point=os.path.join(lr_path, 'linear_regression.py'), role='SageMakerRole',
               train_instance_count=1, train_instance_type=local_instance_type,
               sagemaker_session=sagemaker_local_session, image_name=docker_image,
               framework_version=framework_version, output_path='file://{}'.format(tmpdir))

    data_path = os.path.join(lr_path, 'data')
    s3_prefix = 'integ-test-data/mxnet-linear-regression'
    train_input = sagemaker_local_session.upload_data(path=os.path.join(data_path, 'training'),
                                                      key_prefix=s3_prefix)
    eval_input = sagemaker_local_session.upload_data(path=os.path.join(data_path, 'evaluation'),
                                                     key_prefix=s3_prefix)

    mx.fit({'training': train_input, 'evaluation': eval_input})

    for directory, files in MODEL_SUCCESS_FILES.items():
        local_mode_utils.assert_output_files_exist(str(tmpdir), directory, files)

def test_keras_training(docker_image, sagemaker_local_session, local_instance_type,
                        framework_version, tmpdir):
    keras_path = os.path.join(RESOURCE_PATH, 'keras')
    script_path = os.path.join(keras_path, 'keras_mnist.py')

    mx = MXNet(entry_point=script_path, role='SageMakerRole', train_instance_count=1,
               train_instance_type=local_instance_type, sagemaker_session=sagemaker_local_session,
               image_name=docker_image, framework_version=framework_version,
               output_path='file://{}'.format(tmpdir))

    train = 'file://{}'.format(os.path.join(keras_path, 'data'))
    mx.fit({'train': train})

    for directory, files in MODEL_SUCCESS_FILES.items():
        local_mode_utils.assert_output_files_exist(str(tmpdir), directory, files)

def test_mxnet_neo(strftime, sagemaker_session, mxnet_version, skip_if_mms_version):
    mx = MXNet(
        entry_point=SCRIPT_PATH,
        role=ROLE,
        sagemaker_session=sagemaker_session,
        train_instance_count=INSTANCE_COUNT,
        train_instance_type=INSTANCE_TYPE,
        framework_version=mxnet_version,
    )

    inputs = "s3://mybucket/train"

    mx.fit(inputs=inputs)

    input_shape = {"data": [100, 1, 28, 28]}
    output_location = "s3://neo-sdk-test"

    compiled_model = mx.compile_model(

def test_local_transform_mxnet(
    sagemaker_local_session, tmpdir, mxnet_full_version, cpu_instance_type
):
    data_path = os.path.join(DATA_DIR, "mxnet_mnist")
    script_path = os.path.join(data_path, "mnist.py")

    mx = MXNet(
        entry_point=script_path,
        role="SageMakerRole",
        train_instance_count=1,
        train_instance_type="local",
        framework_version=mxnet_full_version,
        sagemaker_session=sagemaker_local_session,
    )

    train_input = mx.sagemaker_session.upload_data(
        path=os.path.join(data_path, "train"), key_prefix="integ-test-data/mxnet_mnist/train"
    )
    test_input = mx.sagemaker_session.upload_data(
        path=os.path.join(data_path, "test"), key_prefix="integ-test-data/mxnet_mnist/test"
    )

    with stopit.ThreadingTimeout(5 * 60, swallow_exc=False):

def test_estimator_script_mode_launch_parameter_server(sagemaker_session):
    mx = MXNet(
        entry_point=SCRIPT_PATH,
        role=ROLE,
        sagemaker_session=sagemaker_session,
        train_instance_count=INSTANCE_COUNT,
        train_instance_type=INSTANCE_TYPE,
        distributions=LAUNCH_PS_DISTRIBUTIONS_DICT,
        framework_version="1.3.0",
    )
    assert mx.hyperparameters().get(MXNet.LAUNCH_PS_ENV_NAME) == "true"

def test_s3_input_mode(sagemaker_session, tuner):
    expected_input_mode = "Pipe"

    script_path = os.path.join(DATA_DIR, "mxnet_mnist", "failure_script.py")
    mxnet = MXNet(
        entry_point=script_path,
        role=ROLE,
        framework_version=FRAMEWORK_VERSION,
        train_instance_count=TRAIN_INSTANCE_COUNT,
        train_instance_type=TRAIN_INSTANCE_TYPE,
        sagemaker_session=sagemaker_session,
    )
    tuner.estimator = mxnet

    tags = [{"Name": "some-tag-without-a-value"}]
    tuner.tags = tags

    hyperparameter_ranges = {
        "num_components": IntegerParameter(2, 4),
        "algorithm_mode": CategoricalParameter(["regular", "randomized"]),
    }

def test_framework_tuning_config(sagemaker_session):
    mxnet_estimator = mxnet.MXNet(
        entry_point="{{ entry_point }}",
        source_dir="{{ source_dir }}",
        py_version="py3",
        framework_version="1.3.0",
        role="{{ role }}",
        train_instance_count=1,
        train_instance_type="ml.m4.xlarge",
        sagemaker_session=sagemaker_session,
        base_job_name="{{ base_job_name }}",
        hyperparameters={"batch_size": 100},
    )

    hyperparameter_ranges = {
        "optimizer": tuner.CategoricalParameter(["sgd", "Adam"]),
        "learning_rate": tuner.ContinuousParameter(0.01, 0.2),
        "num_epoch": tuner.IntegerParameter(10, 50),

def _create_model(output_path):
        script_path = os.path.join(DATA_DIR, "mxnet_mnist", "mnist.py")
        data_path = os.path.join(DATA_DIR, "mxnet_mnist")

        mx = MXNet(
            entry_point=script_path,
            role="SageMakerRole",
            train_instance_count=1,
            train_instance_type="local",
            output_path=output_path,
            framework_version=mxnet_full_version,
            sagemaker_session=sagemaker_local_session,
        )

        train_input = mx.sagemaker_session.upload_data(
            path=os.path.join(data_path, "train"), key_prefix="integ-test-data/mxnet_mnist/train"
        )
        test_input = mx.sagemaker_session.upload_data(
            path=os.path.join(data_path, "test"), key_prefix="integ-test-data/mxnet_mnist/test"
        )