def lambda_handler()

in custom_resource/sagemaker_query_training.py [0:0]


def lambda_handler(event, context):
    if "TrainingJobName" in event:
        job_name = event["TrainingJobName"]
    else:
        raise KeyError("TrainingJobName not found for event: {}.".format(json.dumps(event)))

    # Get the training job
    response = sm_client.describe_training_job(TrainingJobName=job_name)
    status = response["TrainingJobStatus"]
    logger.info("Training job:{} has status:{}.".format(job_name, status))

    # Get the metrics as a dictionary
    for _, metric in enumerate(response["FinalMetricDataList"]):
        metric["Timestamp"] = metric["Timestamp"].timestamp()

    return {
        "statusCode": 200,
        "results": {
            "TrainingJobName": job_name,
            "TrainingJobStatus": status,
            "TrainingMetrics": response["FinalMetricDataList"],
        },
    }