Software Development Engineer

Blog PostsResume

Restrict Users to Create SageMaker Endpoints from Model with Specific Execution Policy

When you deploy a SageMaker Endpoint, the following operations occur at the backend.

  1. An ML compute instance is provisioned in a service managed account
  2. The Model image is downloaded on the instance and the container is run.
  3. The ML compute instance assumes the execution role specified in Model to perform the requisite operations.
  4. Thereafter, requests are made to the /ping route for a specific period of time before the Sagemaker Endpoint is rendered InService

With the current architecture, a user can deploy an Endpoint using any of the available models and there is no way to implicitly restrict any user from from doing so. Thus, to achieve this use case (rather a workaround to enforce the scenario), you may consider the following procedure.

  1. Create a Cloudwatch Events Rule with Service Name as Sagemaker, Event Type as AWS API Call via CloudTrail and Operation as CreateEndpoint (or CreateTransformJob)
  2. As a Target to this Event Rule, add a Lambda function that does the following:
  • The lambda function stores the Model names and the roles that can access them as constants.

  • Then at each call to this function, it checks whether the API action was called by any of the allowed Roles. If not, then it Stops the resource in concern.

    const AWS = require('aws-sdk');

    exports.handler = (event, context, callback) => { const sagemaker = new AWS.SageMaker(); const roleARN = event["userIdentity"]["sessionContext"]["sessionIssuer"]["arn"]; const modelName = event["requestParameters"]["modelName"]; const transformJob = event["requestParameters"]["transformJobName"]; let roles = process.env.ALLOWED_ROLES.split(","); if (!roles.includes(roleARN)) { var params = { TransformJobName: transformJob };

          sagemaker.stopTransformJob(params, function(err, data) {
              if (err) console.log(err, err.stack); 
              else     console.log(data);  
          });
      }
      callback(null, {});

    };

With this, as soon as an API call is made to Create a Batch Transform job, if that action is not called by any of the allowed roles, it will stop the job immediately.


© 2024 Ujjwal Bhardwaj. All Rights Reserved.