def test_sagemaker_exception_in_deploy_model(self):
mock_client = MagicMock()
mock_exception = ClientError({"Error": {
"Message": "SageMaker broke"
}}, "deploy_model")
mock_client.create_endpoint_config.side_effect = mock_exception
mock_args = self.parser.parse_args(required_args)
with self.assertRaises(Exception):
_utils.deploy_model(mock_client, vars(mock_args))
def main(argv=None):
parser = create_parser()
args = parser.parse_args(argv)
logging.getLogger().setLevel(logging.INFO)
client = _utils.get_sagemaker_client(args.region, args.endpoint_url, assume_role_arn=args.assume_role)
endpoint_name = None
old_endpoint_config_name = None
if (args.update_endpoint and _utils.endpoint_name_exists(client, args.endpoint_name)):
## Get the old endpoint config to cleanup later
old_endpoint_config_name = _utils.get_endpoint_config(client, args.endpoint_name)
logging.info('Submitting Update Endpoint request to SageMaker...')
endpoint_name = _utils.update_deployed_model(client, vars(args))
else:
logging.info('Submitting Create Endpoint request to SageMaker...')
endpoint_name = _utils.deploy_model(client, vars(args))
logging.info('Endpoint creation/update request submitted. Waiting for completion...')
_utils.wait_for_endpoint_creation(client, endpoint_name)
## If updating existing endpoint, cleanup old endpoint config
if old_endpoint_config_name:
logging.info("Deleting old endpoint config: " + old_endpoint_config_name)
if _utils.delete_endpoint_config(client, old_endpoint_config_name):
logging.info("Deleted old endpoint config: " + old_endpoint_config_name)
else:
logging.info("Unable to delete old endpoint config: " + old_endpoint_config_name)
_utils.write_output(args.endpoint_name_output_path, endpoint_name)
logging.info('Endpoint creation/update completed.')
def test_get_endpoint_name_from_job(self):
mock_client = MagicMock()
# if we don't pass --endpoint_name argument then endpoint name is constructed using --model_name_1
self.assertEqual(
_utils.deploy_model(mock_client,
vars(self.parser.parse_args(required_args))),
'Endpoint-test')
def main(argv=None):
parser = create_parser()
args = parser.parse_args(argv)
logging.getLogger().setLevel(logging.INFO)
client = _utils.get_sagemaker_client(args.region, args.endpoint_url)
logging.info('Submitting Endpoint request to SageMaker...')
endpoint_name = _utils.deploy_model(client, vars(args))
logging.info('Endpoint creation request submitted. Waiting for completion...')
_utils.wait_for_endpoint_creation(client, endpoint_name)
with open('/tmp/endpoint_name.txt', 'w') as f:
f.write(endpoint_name)
logging.info('Endpoint creation completed.')
def main(argv=None):
parser = create_parser()
args = parser.parse_args(argv)
logging.getLogger().setLevel(logging.INFO)
client = _utils.get_sagemaker_client(args.region,
args.endpoint_url,
assume_role_arn=args.assume_role)
logging.info('Submitting Endpoint request to SageMaker...')
endpoint_name = _utils.deploy_model(client, vars(args))
logging.info(
'Endpoint creation request submitted. Waiting for completion...')
_utils.wait_for_endpoint_creation(client, endpoint_name)
_utils.write_output(args.endpoint_name_output_path, endpoint_name)
logging.info('Endpoint creation completed.')
def main(argv=None):
parser = argparse.ArgumentParser(description='SageMaker Training Job')
parser.add_argument('--region',
type=str,
help='The region where the cluster launches.')
parser.add_argument('--model_name',
type=str,
help='The name of the new model.')
args = parser.parse_args()
logging.getLogger().setLevel(logging.INFO)
client = _utils.get_client(args.region)
logging.info('Submitting Endpoint request to SageMaker...')
endpoint_name = _utils.deploy_model(client, args.model_name)
with open('/tmp/endpoint_name.txt', 'w') as f:
f.write(endpoint_name)
logging.info('Endpoint creation completed.')
def test_deploy_model(self):
mock_client = MagicMock()
mock_args = self.parser.parse_args(required_args + [
'--endpoint_name', 'test-endpoint-name', '--endpoint_config_name',
'test-endpoint-config-name'
])
response = _utils.deploy_model(mock_client, vars(mock_args))
mock_client.create_endpoint_config.assert_called_once_with(
EndpointConfigName='test-endpoint-config-name',
ProductionVariants=[{
'VariantName': 'variant-name-1',
'ModelName': 'model-test',
'InitialInstanceCount': 1,
'InstanceType': 'ml.m4.xlarge',
'InitialVariantWeight': 1.0
}],
Tags=[])
self.assertEqual(response, 'test-endpoint-name')
def test_get_endpoint_name_from_job(self):
mock_client = MagicMock()
# if we don't pass --endpoint_name argument then endpoint name is constructed
self.assertTrue('Endpoint-' in _utils.deploy_model(
mock_client, vars(self.parser.parse_args(required_args))))
def main(argv=None):
parser = argparse.ArgumentParser(description='SageMaker Training Job')
parser.add_argument('--region',
type=str.strip,
required=True,
help='The region where the cluster launches.')
parser.add_argument('--endpoint_config_name',
type=str.strip,
required=False,
help='The name of the endpoint configuration.',
default='')
parser.add_argument('--variant_name_1',
type=str.strip,
required=False,
help='The name of the production variant.',
default='variant-name-1')
parser.add_argument('--model_name_1',
type=str.strip,
required=True,
help='The model name used for endpoint deployment.')
parser.add_argument('--initial_instance_count_1',
type=_utils.str_to_int,
required=False,
help='Number of instances to launch initially.',
default=1)
parser.add_argument(
'--instance_type_1',
choices=[
'ml.m4.xlarge', 'ml.m4.2xlarge', 'ml.m4.4xlarge', 'ml.m4.10xlarge',
'ml.m4.16xlarge', 'ml.m5.large', 'ml.m5.xlarge', 'ml.m5.2xlarge',
'ml.m5.4xlarge', 'ml.m5.12xlarge', 'ml.m5.24xlarge',
'ml.c4.xlarge', 'ml.c4.2xlarge', 'ml.c4.4xlarge', 'ml.c4.8xlarge',
'ml.p2.xlarge', 'ml.p2.8xlarge', 'ml.p2.16xlarge', 'ml.p3.2xlarge',
'ml.p3.8xlarge', 'ml.p3.16xlarge', 'ml.c5.xlarge', 'ml.c5.2xlarge',
'ml.c5.4xlarge', 'ml.c5.9xlarge', 'ml.c5.18xlarge', ''
],
type=str.strip,
required=False,
help='The ML compute instance type.',
default='ml.m4.xlarge')
parser.add_argument(
'--initial_variant_weight_1',
type=_utils.str_to_float,
required=False,
help=
'Determines initial traffic distribution among all of the models that you specify in the endpoint configuration.',
default=1.0)
parser.add_argument(
'--accelerator_type_1',
choices=['ml.eia1.medium', 'ml.eia1.large', 'ml.eia1.xlarge', ''],
type=str.strip,
required=False,
help=
'The size of the Elastic Inference (EI) instance to use for the production variant.',
default='')
parser.add_argument('--variant_name_2',
type=str.strip,
required=False,
help='The name of the production variant.',
default='variant-name-2')
parser.add_argument('--model_name_2',
type=str.strip,
required=False,
help='The model name used for endpoint deployment.',
default='')
parser.add_argument('--initial_instance_count_2',
type=_utils.str_to_int,
required=False,
help='Number of instances to launch initially.',
default=1)
parser.add_argument(
'--instance_type_2',
choices=[
'ml.m4.xlarge', 'ml.m4.2xlarge', 'ml.m4.4xlarge', 'ml.m4.10xlarge',
'ml.m4.16xlarge', 'ml.m5.large', 'ml.m5.xlarge', 'ml.m5.2xlarge',
'ml.m5.4xlarge', 'ml.m5.12xlarge', 'ml.m5.24xlarge',
'ml.c4.xlarge', 'ml.c4.2xlarge', 'ml.c4.4xlarge', 'ml.c4.8xlarge',
'ml.p2.xlarge', 'ml.p2.8xlarge', 'ml.p2.16xlarge', 'ml.p3.2xlarge',
'ml.p3.8xlarge', 'ml.p3.16xlarge', 'ml.c5.xlarge', 'ml.c5.2xlarge',
'ml.c5.4xlarge', 'ml.c5.9xlarge', 'ml.c5.18xlarge', ''
],
type=str.strip,
required=False,
help='The ML compute instance type.',
default='ml.m4.xlarge')
parser.add_argument(
'--initial_variant_weight_2',
type=_utils.str_to_float,
required=False,
help=
'Determines initial traffic distribution among all of the models that you specify in the endpoint configuration.',
default=1.0)
parser.add_argument(
'--accelerator_type_2',
choices=['ml.eia1.medium', 'ml.eia1.large', 'ml.eia1.xlarge', ''],
type=str.strip,
required=False,
help=
'The size of the Elastic Inference (EI) instance to use for the production variant.',
default='')
parser.add_argument('--variant_name_3',
type=str.strip,
required=False,
help='The name of the production variant.',
default='variant-name-3')
parser.add_argument('--model_name_3',
type=str.strip,
required=False,
help='The model name used for endpoint deployment.',
default='')
parser.add_argument('--initial_instance_count_3',
type=_utils.str_to_int,
required=False,
help='Number of instances to launch initially.',
default=1)
parser.add_argument(
'--instance_type_3',
choices=[
'ml.m4.xlarge', 'ml.m4.2xlarge', 'ml.m4.4xlarge', 'ml.m4.10xlarge',
'ml.m4.16xlarge', 'ml.m5.large', 'ml.m5.xlarge', 'ml.m5.2xlarge',
'ml.m5.4xlarge', 'ml.m5.12xlarge', 'ml.m5.24xlarge',
'ml.c4.xlarge', 'ml.c4.2xlarge', 'ml.c4.4xlarge', 'ml.c4.8xlarge',
'ml.p2.xlarge', 'ml.p2.8xlarge', 'ml.p2.16xlarge', 'ml.p3.2xlarge',
'ml.p3.8xlarge', 'ml.p3.16xlarge', 'ml.c5.xlarge', 'ml.c5.2xlarge',
'ml.c5.4xlarge', 'ml.c5.9xlarge', 'ml.c5.18xlarge', ''
],
type=str.strip,
required=False,
help='The ML compute instance type.',
default='ml.m4.xlarge')
parser.add_argument(
'--initial_variant_weight_3',
type=_utils.str_to_float,
required=False,
help=
'Determines initial traffic distribution among all of the models that you specify in the endpoint configuration.',
default=1.0)
parser.add_argument(
'--accelerator_type_3',
choices=['ml.eia1.medium', 'ml.eia1.large', 'ml.eia1.xlarge', ''],
type=str.strip,
required=False,
help=
'The size of the Elastic Inference (EI) instance to use for the production variant.',
default='')
parser.add_argument(
'--resource_encryption_key',
type=str.strip,
required=False,
help=
'The AWS KMS key that Amazon SageMaker uses to encrypt data on the storage volume attached to the ML compute instance(s).',
default='')
parser.add_argument(
'--endpoint_config_tags',
type=_utils.str_to_json_dict,
required=False,
help='An array of key-value pairs, to categorize AWS resources.',
default='{}')
parser.add_argument('--endpoint_name',
type=str.strip,
required=False,
help='The name of the endpoint.',
default='')
parser.add_argument(
'--endpoint_tags',
type=_utils.str_to_json_dict,
required=False,
help='An array of key-value pairs, to categorize AWS resources.',
default='{}')
args = parser.parse_args()
logging.getLogger().setLevel(logging.INFO)
client = _utils.get_client(args.region)
logging.info('Submitting Endpoint request to SageMaker...')
endpoint_name = _utils.deploy_model(client, vars(args))
logging.info(
'Endpoint creation request submitted. Waiting for completion...')
_utils.wait_for_endpoint_creation(client, endpoint_name)
with open('/tmp/endpoint_name.txt', 'w') as f:
f.write(endpoint_name)
logging.info('Endpoint creation completed.')
