feat(apigateway): L2 construct for Sagemaker Integration (#25459)

Sagemaker integration for ApiGateway. Pulled the IEndpoint from the alpha module. Stabilizes `sagemaker.IEndpoint` from `aws-sagemaker-alpha` so that it can be used in the integration. ---- *By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
aws · Aug 25, 2023 · 53d61bb · 53d61bb
1 parent b5bd39e
commit 53d61bb
Show file tree

Hide file tree

Showing 10 changed files with 169 additions and 35 deletions.
diff --git a/packages/@aws-cdk/aws-sagemaker-alpha/lib/endpoint.ts b/packages/@aws-cdk/aws-sagemaker-alpha/lib/endpoint.ts
@@ -4,6 +4,7 @@ import * as cloudwatch from 'aws-cdk-lib/aws-cloudwatch';
 import * as ec2 from 'aws-cdk-lib/aws-ec2';
 import * as iam from 'aws-cdk-lib/aws-iam';
 import * as cdk from 'aws-cdk-lib/core';
+import * as sagemaker from 'aws-cdk-lib/aws-sagemaker';
 import { Construct } from 'constructs';
 import { EndpointConfig, IEndpointConfig, InstanceProductionVariant } from './endpoint-config';
 import { InstanceType } from './instance-type';
@@ -20,30 +21,12 @@ const BURSTABLE_INSTANCE_TYPE_PREFIXES = Object.entries(ec2.InstanceClass)
   .filter(([name, _]) => name.startsWith('T'))
   .map(([_, prefix]) => `ml.${prefix}.`);
 
+// IEndpoint is stabilized so that it can be used in aws-apigateway SagemakerIntegration
+// Exposing it again here so that there is no breakage to aws-sagemaker-alpha
 /**
- * The interface for a SageMaker Endpoint resource.
+ * The Interface for a SageMaker Endpoint resource.
  */
-export interface IEndpoint extends cdk.IResource {
-  /**
-   * The ARN of the endpoint.
-   *
-   * @attribute
-   */
-  readonly endpointArn: string;
-
-  /**
-   * The name of the endpoint.
-   *
-   * @attribute
-   */
-  readonly endpointName: string;
-
-  /**
-   * Permits an IAM principal to invoke this endpoint
-   * @param grantee The principal to grant access to
-   */
-  grantInvoke(grantee: iam.IGrantable): iam.Grant;
-}
+export interface IEndpoint extends sagemaker.IEndpoint {}
 
 /**
  * Represents the features common to all production variant types (e.g., instance, serverless) that

diff --git a/packages/aws-cdk-lib/aws-apigateway/README.md b/packages/aws-cdk-lib/aws-apigateway/README.md
@@ -257,9 +257,10 @@ method is called. API Gateway supports the following integrations:
 
 - `MockIntegration` - can be used to test APIs. This is the default
    integration if one is not specified.
-- `LambdaIntegration` - can be used to invoke an AWS Lambda function.
 - `AwsIntegration` - can be used to invoke arbitrary AWS service APIs.
 - `HttpIntegration` - can be used to invoke HTTP endpoints.
+- `LambdaIntegration` - can be used to invoke an AWS Lambda function.
+- `SagemakerIntegration` - can be used to invoke Sagemaker Endpoints.
 
 The following example shows how to integrate the `GET /book/{book_id}` method to
 an AWS Lambda function:

diff --git a/packages/aws-cdk-lib/aws-apigateway/lib/integrations/index.ts b/packages/aws-cdk-lib/aws-apigateway/lib/integrations/index.ts
@@ -3,4 +3,5 @@ export * from './lambda';
 export * from './http';
 export * from './mock';
 export * from './stepfunctions';
+export * from './sagemaker';
 export * from './request-context';
diff --git a/packages/aws-cdk-lib/aws-apigateway/lib/integrations/sagemaker.ts b/packages/aws-cdk-lib/aws-apigateway/lib/integrations/sagemaker.ts
@@ -0,0 +1,64 @@
+import { AwsIntegration } from './aws';
+import * as iam from '../../../aws-iam';
+import { IEndpoint } from '../../../aws-sagemaker';
+import { IntegrationConfig, IntegrationOptions } from '../integration';
+import { Method } from '../method';
+
+/**
+ * Options for SageMakerIntegration
+ */
+export interface SagemakerIntegrationOptions extends IntegrationOptions {
+}
+
+/**
+ * Integrates an AWS Sagemaker Endpoint to an API Gateway method
+ *
+ * @example
+ *
+ *   declare const resource: apigateway.Resource;
+ *   declare const endpoint: sagemaker.IEndpoint;
+ *   resource.addMethod('POST', new apigateway.SagemakerIntegration(endpoint));
+ *
+ */
+export class SagemakerIntegration extends AwsIntegration {
+  private readonly endpoint: IEndpoint;
+
+  constructor(endpoint: IEndpoint, options: SagemakerIntegrationOptions = {}) {
+    super({
+      service: 'runtime.sagemaker',
+      path: `endpoints/${endpoint.endpointName}/invocations`,
+      options: {
+        credentialsRole: options.credentialsRole,
+        integrationResponses: [
+          { statusCode: '200' },
+        ],
+        ...options,
+      },
+    });
+
+    this.endpoint = endpoint;
+  }
+
+  public bind(method: Method): IntegrationConfig {
+    const bindResult = super.bind(method);
+
+    const credentialsRole = bindResult.options?.credentialsRole ?? new iam.Role(method, 'SagemakerRole', {
+      assumedBy: new iam.ServicePrincipal('apigateway.amazonaws.com'),
+      description: 'Generated by CDK::ApiGateway::SagemakerIntegration',
+    });
+
+    this.endpoint.grantInvoke(credentialsRole);
+
+    method.addMethodResponse({
+      statusCode: '200',
+    });
+
+    return {
+      ...bindResult,
+      options: {
+        ...bindResult.options,
+        credentialsRole,
+      },
+    };
+  }
+}
diff --git a/packages/aws-cdk-lib/aws-apigateway/lib/integrations/stepfunctions.ts b/packages/aws-cdk-lib/aws-apigateway/lib/integrations/stepfunctions.ts
@@ -8,11 +8,11 @@ import { Token } from '../../../core';
 import { IntegrationConfig, IntegrationOptions, PassthroughBehavior } from '../integration';
 import { Method } from '../method';
 import { Model } from '../model';
+
 /**
  * Options when configuring Step Functions synchronous integration with Rest API
  */
 export interface StepFunctionsExecutionIntegrationOptions extends IntegrationOptions {
-
   /**
    * Which details of the incoming request must be passed onto the underlying state machine,
    * such as, account id, user identity, request id, etc. The execution input will include a new key `requestContext`:

diff --git a/packages/aws-cdk-lib/aws-apigateway/test/integrations/sagemaker.test.ts b/packages/aws-cdk-lib/aws-apigateway/test/integrations/sagemaker.test.ts
@@ -0,0 +1,56 @@
+import { Template } from '../../../assertions';
+import * as iam from '../../../aws-iam';
+import * as sagemaker from '../../../aws-sagemaker';
+import * as cdk from '../../../core';
+import * as apigateway from '../../lib';
+
+describe('SageMaker Integration', () => {
+  test('minimal setup', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+    const api = new apigateway.RestApi(stack, 'my-api');
+    const endpoint = new FakeEndpoint(stack, 'FakeEndpoint');
+
+    // WHEN
+    const integration = new apigateway.SagemakerIntegration(endpoint);
+    api.root.addMethod('POST', integration);
+
+    // THEN
+    Template.fromStack(stack).hasResourceProperties('AWS::ApiGateway::Method', {
+      Integration: {
+        IntegrationHttpMethod: 'POST',
+        Type: 'AWS',
+        Uri: {
+          'Fn::Join': [
+            '',
+            [
+              'arn:',
+              {
+                Ref: 'AWS::Partition',
+              },
+              ':apigateway:',
+              {
+                Ref: 'AWS::Region',
+              },
+              ':runtime.sagemaker:path/endpoints/endpointName/invocations',
+            ],
+          ],
+        },
+      },
+    });
+  });
+});
+
+class FakeEndpoint extends cdk.Resource implements sagemaker.IEndpoint {
+  public readonly endpointArn = 'endpointArn';
+
+  public readonly endpointName = 'endpointName';
+
+  public grantInvoke(grantee: iam.IGrantable) {
+    return iam.Grant.addToPrincipal({
+      grantee,
+      actions: ['sagemaker:InvokeEndpoint'],
+      resourceArns: [this.endpointArn],
+    });
+  }
+}
diff --git a/packages/aws-cdk-lib/aws-sagemaker/lib/endpoint.ts b/packages/aws-cdk-lib/aws-sagemaker/lib/endpoint.ts
@@ -0,0 +1,27 @@
+import * as iam from '../../aws-iam';
+import { IResource } from '../../core';
+
+/**
+ * The interface for a SageMaker Endpoint resource.
+ */
+export interface IEndpoint extends IResource {
+  /**
+   * The ARN of the endpoint.
+   *
+   * @attribute
+   */
+  readonly endpointArn: string;
+
+  /**
+   * The name of the endpoint.
+   *
+   * @attribute
+   */
+  readonly endpointName: string;
+
+  /**
+   * Permits an IAM principal to invoke this endpoint
+   * @param grantee The principal to grant access to
+   */
+  grantInvoke(grantee: iam.IGrantable): iam.Grant;
+}
diff --git a/packages/aws-cdk-lib/aws-sagemaker/lib/index.ts b/packages/aws-cdk-lib/aws-sagemaker/lib/index.ts
@@ -1 +1,2 @@
-export * from './sagemaker.generated';
+export * from './sagemaker.generated';
+export * from './endpoint';
diff --git a/packages/aws-cdk-lib/rosetta/aws_apigateway/default.ts-fixture b/packages/aws-cdk-lib/rosetta/aws_apigateway/default.ts-fixture
@@ -1,14 +1,15 @@
 // Fixture with packages imported, but nothing else
 import { Construct } from 'constructs';
 import { Duration, Stack } from 'aws-cdk-lib';
-import apigateway = require('aws-cdk-lib/aws-apigateway');
-import cognito = require('aws-cdk-lib/aws-cognito');
-import lambda = require('aws-cdk-lib/aws-lambda');
-import iam = require('aws-cdk-lib/aws-iam');
-import s3 = require('aws-cdk-lib/aws-s3');
-import ec2 = require('aws-cdk-lib/aws-ec2');
-import logs = require('aws-cdk-lib/aws-logs');
-import stepfunctions = require('aws-cdk-lib/aws-stepfunctions');
+import * as apigateway from 'aws-cdk-lib/aws-apigateway';
+import * as cognito from 'aws-cdk-lib/aws-cognito';
+import * as lambda from 'aws-cdk-lib/aws-lambda';
+import * as iam from 'aws-cdk-lib/aws-iam';
+import * as s3 from 'aws-cdk-lib/aws-s3';
+import * as ec2 from 'aws-cdk-lib/aws-ec2';
+import * as logs from 'aws-cdk-lib/aws-logs';
+import * as stepfunctions from 'aws-cdk-lib/aws-stepfunctions';
+import * as sagemaker from 'aws-cdk-lib/aws-sagemaker';
 
 class Fixture extends Stack {
   constructor(scope: Construct, id: string) {

diff --git a/packages/aws-cdk-lib/rosetta/aws_apigateway/stepfunctions.ts-fixture b/packages/aws-cdk-lib/rosetta/aws_apigateway/stepfunctions.ts-fixture
@@ -1,7 +1,7 @@
 import { Construct } from 'constructs';
 import { Stack } from 'aws-cdk-lib';
-import apigateway = require('aws-cdk-lib/aws-apigateway');
-import stepfunctions = require('aws-cdk-lib/aws-stepfunctions');
+import * as apigateway from 'aws-cdk-lib/aws-apigateway';
+import * as stepfunctions from 'aws-cdk-lib/aws-stepfunctions';
 
 class Fixture extends Stack {
   constructor(scope: Construct, id: string) {