def setup_agent()

in setup/run.py [0:0]

• 52 lines of code
• 5 McCabe index (conditional complexity)

def setup_agent(agent_id, args, thing_group_name, thing_group_arn):
    policy_name='WindTurbineFarmPolicy-%s' % args.sagemaker_project_id
    base="agent/certificates/iot/edge_device_%d_%s.pem"
    fleet_name = 'wind-turbine-farm-%s' % args.sagemaker_project_id
    thing_arn_template = thing_group_arn.replace('thinggroup', 'thing').replace(thing_group_name, '%s')
    cred_host = iot_client.describe_endpoint(endpointType='iot:CredentialProvider')['endpointAddress']
    policy_alias = 'SageMakerEdge-%s' % fleet_name

    # register the device in the fleet    
    # the device name needs to have 36 chars
    dev_name = "edge-device-%d" % agent_id
    dev = [{'DeviceName': dev_name, 'IotThingName': dev_name}]    
    try:        
        sm_client.describe_device(DeviceFleetName=fleet_name, DeviceName=dev_name)
        logger.info("Device was already registered on SageMaker Edge Manager")
    except ClientError as e:
        if e.response['Error']['Code'] != 'ValidationException': raise e
        logger.info("Registering a new device %s on fleet %s" % (dev_name, fleet_name))
        sm_client.register_devices(DeviceFleetName=fleet_name, Devices=dev)
        iot_client.add_thing_to_thing_group(
            thingGroupName=thing_group_name,
            thingGroupArn=thing_group_arn,
            thingName='edge-device-%d' % agent_id,
            thingArn=thing_arn_template % ('edge-device-%d' % agent_id)
        )        
    
    # if you reach this point you need to create new certificates
    # generate the certificates    
    cert=base % (agent_id, 'cert')
    key=base % (agent_id, 'pub')
    pub=base % (agent_id, 'key')
    
    cert_meta=iot_client.create_keys_and_certificate(setAsActive=True)
    cert_arn = cert_meta['certificateArn']
    with open(cert, 'w') as c: c.write(cert_meta['certificatePem'])
    with open(key,  'w') as c: c.write(cert_meta['keyPair']['PrivateKey'])
    with open(pub,  'w') as c: c.write(cert_meta['keyPair']['PublicKey'])
        
    # attach the certificates to the policy and to the thing
    iot_client.attach_policy(policyName=policy_name, target=cert_arn)
    iot_client.attach_thing_principal(thingName='edge-device-%d' % agent_id, principal=cert_arn)        
    
    logger.info("Finally, let's create the agent config file")
    agent_params = {
        "sagemaker_edge_core_device_name": dev_name,
        "sagemaker_edge_core_device_fleet_name": fleet_name,
        "sagemaker_edge_core_capture_data_buffer_size": 30,
        "sagemaker_edge_core_capture_data_batch_size": 10,
        "sagemaker_edge_core_capture_data_push_period_seconds": 4,
        "sagemaker_edge_core_folder_prefix": "wind_turbine_data",
        "sagemaker_edge_core_region": args.aws_region,
        "sagemaker_edge_core_root_certs_path": "./agent/certificates/root",
        "sagemaker_edge_provider_aws_ca_cert_file":"./agent/certificates/iot/AmazonRootCA1.pem",
        "sagemaker_edge_provider_aws_cert_file":"./%s" % cert,
        "sagemaker_edge_provider_aws_cert_pk_file":"./%s" % key,
        "sagemaker_edge_provider_aws_iot_cred_endpoint": "https://%s/role-aliases/%s/credentials" % (cred_host,policy_alias),
        "sagemaker_edge_provider_provider": "Aws",
        "sagemaker_edge_provider_s3_bucket_name": args.artifact_bucket,
        "sagemaker_edge_core_capture_data_destination": "Cloud"
    }
    with open('agent/conf/config_edge_device_%d.json' % agent_id, 'w') as conf:
        conf.write(json.dumps(agent_params, indent=4))