def test_model_image(sagemaker_session): knn = KNN(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM, channel='train') knn.fit(data, MINI_BATCH_SIZE) model = knn.create_model() assert model.image == registry(REGION, "knn") + '/knn:1'
def test_predictor_type(sagemaker_session): knn = KNN(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM, channel='train') knn.fit(data, MINI_BATCH_SIZE) model = knn.create_model() predictor = model.deploy(1, TRAIN_INSTANCE_TYPE) assert isinstance(predictor, KNNPredictor)
def test_call_fit_none_mini_batch_size(sagemaker_session): knn = KNN(base_job_name="knn", sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM, channel='train') knn.fit(data)
def test_call_fit(base_fit, sagemaker_session): knn = KNN(base_job_name="knn", sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM, channel='train') knn.fit(data, MINI_BATCH_SIZE) base_fit.assert_called_once() assert len(base_fit.call_args[0]) == 2 assert base_fit.call_args[0][0] == data assert base_fit.call_args[0][1] == MINI_BATCH_SIZE
def test_model_image(sagemaker_session): knn = KNN(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) data = RecordSet( "s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM, channel="train", ) knn.fit(data, MINI_BATCH_SIZE) model = knn.create_model() assert image_uris.retrieve("knn", REGION) == model.image_uri
def test_predictor_custom_serialization(sagemaker_session): knn = KNN(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) data = RecordSet( "s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM, channel="train", ) knn.fit(data, MINI_BATCH_SIZE) model = knn.create_model() custom_serializer = Mock() custom_deserializer = Mock() predictor = model.deploy( 1, INSTANCE_TYPE, serializer=custom_serializer, deserializer=custom_deserializer, ) assert isinstance(predictor, KNNPredictor) assert predictor.serializer is custom_serializer assert predictor.deserializer is custom_deserializer