def test_sagemaker_exception_in_workteam(self):
mock_client = MagicMock()
mock_exception = ClientError({"Error": {"Message": "SageMaker broke"}}, "workteam")
mock_client.create_workteam.side_effect = mock_exception
mock_args = self.parser.parse_args(required_args)
with self.assertRaises(Exception):
_utils.create_workteam(mock_client, vars(mock_args))
def main(argv=None):
parser = argparse.ArgumentParser(
description='SageMaker Hyperparameter Tuning Job')
parser.add_argument('--region',
type=str.strip,
required=True,
help='The region where the cluster launches.')
parser.add_argument('--team_name',
type=str.strip,
required=True,
help='The name of your work team.')
parser.add_argument('--description',
type=str.strip,
required=True,
help='A description of the work team.')
parser.add_argument(
'--user_pool',
type=str.strip,
required=False,
help=
'An identifier for a user pool. The user pool must be in the same region as the service that you are calling.',
default='')
parser.add_argument(
'--user_groups',
type=str.strip,
required=False,
help='A list of identifiers for user groups separated by commas.',
default='')
parser.add_argument(
'--client_id',
type=str.strip,
required=False,
help=
'An identifier for an application client. You must create the app client ID using Amazon Cognito.',
default='')
parser.add_argument(
'--sns_topic',
type=str.strip,
required=False,
help=
'The ARN for the SNS topic to which notifications should be published.',
default='')
parser.add_argument(
'--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 a create workteam request to SageMaker...')
workteam_arn = _utils.create_workteam(client, vars(args))
logging.info('Workteam created.')
with open('/tmp/workteam_arn.txt', 'w') as f:
f.write(workteam_arn)
def test_get_workteam_output_from_job(self):
mock_client = MagicMock()
mock_client.create_workteam.return_value = {"WorkteamArn": "fake-arn"}
self.assertEqual(
_utils.create_workteam(mock_client,
vars(
self.parser.parse_args(required_args))),
'fake-arn')
def test_workteam(self):
mock_client = MagicMock()
mock_args = self.parser.parse_args(required_args)
response = _utils.create_workteam(mock_client, vars(mock_args))
mock_client.create_workteam.assert_called_once_with(
Description='fake team',
MemberDefinitions=[{'CognitoMemberDefinition': {'UserPool': '', 'UserGroup': '', 'ClientId': ''}}], Tags=[],
WorkteamName='test-team'
)
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 a create workteam request to SageMaker...')
workteam_arn = _utils.create_workteam(client, vars(args))
logging.info('Workteam created.')
_utils.write_output(args.workteam_arn_output_path, workteam_arn)
def main(argv=None):
parser = create_parser()
args = parser.parse_args()
logging.getLogger().setLevel(logging.INFO)
client = _utils.get_sagemaker_client(args.region, args.endpoint_url)
logging.info('Submitting a create workteam request to SageMaker...')
workteam_arn = _utils.create_workteam(client, vars(args))
logging.info('Workteam created.')
with open('/tmp/workteam_arn.txt', 'w') as f:
f.write(workteam_arn)
