Monitor SageMaker Endpoint Status and Trigger Lambda Function

Monday, January 21, 2019 by Ujjwal Bhardwaj

Amazon CloudWatch Events enables you to react selectively to events in the cloud as well as in your applications. Using simple rules you can easily route each type of event to one or more targets including but not limited to AWS Lambda functions, Amazon SNS topics etc.

In this blog post, we will monitor that status of SageMaker Endpoint and then Trigger a Lambda Function to perform operations when the Endpoint gets "InService".

Create a Lambda function pollSageMakerEndpoint to poll the status of Sagemaker endpoint and perform the requisite operation.
Create a Cloudwatch Event Rule to set a Schedule to Invoke the Lambda function pollSageMakerEndpoint.
Create a Lambda function enableCloudwatchRule to Enable the Cloudwatch function that polls Sagemaker endpoing on schedule.
Create a Cloudwatch Events Rule to invoke the function enableCloudwatchRule when SageMaker endpoint is Created/Updated.

pollSageMakerEndpoint

const AWS = require('aws-sdk')

exports.handler = (event, context, callback) => {
  var sagemaker = new AWS.SageMaker();
  var lambda = new AWS.Lambda();
  var cloudwatchevents = new AWS.CloudWatchEvents();

  var params = {
    EndpointName: process.env.ENDPOINT_NAME
  };

  sagemaker.describeEndpoint(params, function(err, data) {
    if (err) {
      console.log(err, err.stack);
    }
    else {
      console.log(data.EndpointStatus);
      if (data.EndpointStatus !== 'InService') {
        var params = {
          FunctionName: process.env.TRIGGER_FUNCTION,
        };
        lambda.invoke(params, function(err, data) {
          if (err) console.log(err, err.stack);
          else { 
            console.log(data);
            var params = {
              Name: process.env.CLOUDWATCH_RULE
            };
            cloudwatchevents.disableRule(params, function(err, data) {
              if (err) console.log(err, err.stack);
              else {
                console.log(data);
                const response = {
                  statusCode: 200,
                  body: JSON.stringify(data),
                };
                callback(null, response);
              }
            });
          }
        });
      }
    }
  });
};

Screen Shot 2019-02-05 at 7.16.25 PM

enableCloudwatchRule

const AWS = require('aws-sdk')

exports.handler = (event, context, callback) => {
  var cloudwatchevents = new AWS.CloudWatchEvents();

  var params = {
    Name: process.env.CLOUDWATCH_RULE
  };

  cloudwatchevents.enableRule(params, function(err, data) {
    if (err) console.log(err, err.stack);
    else {
      console.log(data);
      const response = {
        statusCode: 200,
        body: JSON.stringify(data),
      };
      callback(null, response);
    }
  });
};

Screen Shot 2019-02-05 at 7.15.36 PM