def get_config()

in tensorflow_script_mode_local_training_and_serving/tensorflow_script_mode_local_training_and_serving.py [0:0]

• 29 lines of code
• 3 McCabe index (conditional complexity)

def get_config(mode):
    assert mode is CLOUD_MODE or mode is LOCAL_MODE, f'unknown mode selected: {mode}'

    if mode == CLOUD_MODE:
        ## REPLACE WITH A VALID IAM ROLE - START ##
        role = DUMMY_IAM_ROLE
        ## REPLACE WITH A VALID IAM ROLE - END ##
        assert role is not DUMMY_IAM_ROLE, "For cloud mode set a valid sagemaker iam role"

        print('Will run training on an ML instance in AWS.')
        session = sagemaker.Session()
        bucket = session.default_bucket()
        s3_data_prefix = 'tensorflow_script_mode_cloud_training/mnist/'
        instance_type = 'ml.m5.large'
        training_dataset_path = 's3://' + bucket + '/' + s3_data_prefix

    else:  # mode == LOCAL_MODE
        print('Will run training locally in a container image.')
        session = LocalSession()
        session.config = {'local': {'local_code': True}}
        instance_type = 'local'
        training_dataset_path = "file://./data/"
        role = DUMMY_IAM_ROLE  # not needed in local training
        s3_data_prefix = None  # not needed in local training
        bucket = None  # not needed in local training

    config = {
        'mode': mode,
        's3_data_prefix': s3_data_prefix,
        'sagemaker_session': session,
        'bucket': bucket,
        'instance_type': instance_type,
        'training_dataset_path': training_dataset_path,
        'role': role}
    return config