feat: add a test multi-container endpoint

Author: jtblack

lib/osml/model_endpoint/me_sm_endpoint.ts

@@ -11,6 +11,33 @@ import { Construct } from "constructs";
 
 import { BaseConfig, ConfigType } from "../utils/base_config";
 
+/**
+ * Configuration for a container in a SageMaker model
+ */
+export interface ContainerDefinition {
+  /**
+   * The URI of the container image
+   */
+  imageUri: string;
+
+  /**
+   * Environment variables for the container
+   */
+  environment?: Record<string, unknown>;
+
+  /**
+   * Repository access mode for the container
+   */
+  repositoryAccessMode?: string;
+
+  /**
+   * Hostname for the container (required for multi-container Direct mode)
+   * If required and not specified, will default to an index-based name
+   * example: "container-0"
+   */
+  containerHostname?: string;
+}
+
 /**
  * Configuration class for MESMEndpoint Construct.
  */
@@ -38,13 +65,20 @@ export class MESMEndpointConfig extends BaseConfig {
    */
   public SECURITY_GROUP_ID: string;
 
+  /**
+   * List of container definitions for the model
+   */
+  public CONTAINERS: ContainerDefinition[];
+
   /**
    * A JSON object which includes ENV variables to be put into the model container.
+   * @deprecated Use CONTAINERS (ContainerDefinition[]) instead
    */
   public CONTAINER_ENV: Record<string, unknown>;
 
   /**
    * The repository access mode to use for the SageMaker endpoint container.
+   * @deprecated Use CONTAINERS (ContainerDefinition) instead
    */
   public REPOSITORY_ACCESS_MODE: string;
   /**
@@ -57,9 +91,30 @@ export class MESMEndpointConfig extends BaseConfig {
       INITIAL_VARIANT_WEIGHT: 1,
       INITIAL_INSTANCE_COUNT: 1,
       VARIANT_NAME: "AllTraffic",
-      REPOSITORY_ACCESS_MODE: "Platform",
+      CONTAINERS: [],
       ...config
     });
+
+    // Convert deprecated interface to container list if needed
+    if (this.CONTAINERS.length === 0 && config.CONTAINER_ENV !== undefined) {
+      this.CONTAINERS = [
+        {
+          imageUri: "", // Populated later with props.containerImageUri
+          environment: config.CONTAINER_ENV as Record<string, unknown>,
+          repositoryAccessMode: (config.REPOSITORY_ACCESS_MODE ||
+            "Platform") as string
+        }
+      ];
+    } else if (this.CONTAINERS.length === 0) {
+      // Ensure we always have a CONTAINERS array - default to an empty container definition
+      this.CONTAINERS = [
+        {
+          imageUri: "",
+          environment: {} as Record<string, unknown>,
+          repositoryAccessMode: "Platform"
+        }
+      ];
+    }
   }
 }
 
@@ -151,26 +206,35 @@ export class MESMEndpoint extends Construct {
       this.config = [props.config];
     }
 
-    const models = this.config.map(
-      (config) =>
-        new CfnModel(this, `${id}-${config.VARIANT_NAME}`, {
-          executionRoleArn: props.roleArn,
-          containers: [
-            {
-              image: props.containerImageUri,
-              environment: config.CONTAINER_ENV,
-              imageConfig: {
-                repositoryAccessMode:
-                  config.REPOSITORY_ACCESS_MODE || "Platform"
-              }
-            }
-          ],
-          vpcConfig: {
-            subnets: props.subnetIds,
-            securityGroupIds: [config.SECURITY_GROUP_ID]
-          }
-        })
-    );
+    const models = this.config.map((config) => {
+      // Set the imageUri for containers that don't have one specified. This
+      //  handles the legacy conversion case where imageUri was initially empty.
+      config.CONTAINERS = config.CONTAINERS.map((container) => ({
+        ...container,
+        imageUri: container.imageUri || props.containerImageUri
+      }));
+
+      // Map to the SageMaker container format
+      const containers = config.CONTAINERS.map((container, index) => ({
+        image: container.imageUri,
+        environment: container.environment || {},
+        imageConfig: {
+          repositoryAccessMode: container.repositoryAccessMode || "Platform"
+        },
+        containerHostname: container.containerHostname || `container-${index}`
+      }));
+
+      return new CfnModel(this, `${id}-${config.VARIANT_NAME}`, {
+        executionRoleArn: props.roleArn,
+        containers: containers,
+        inferenceExecutionConfig:
+          containers.length > 1 ? { mode: "Direct" } : undefined,
+        vpcConfig: {
+          subnets: props.subnetIds,
+          securityGroupIds: [config.SECURITY_GROUP_ID]
+        }
+      });
+    });
 
     this.endpointConfig = new CfnEndpointConfig(this, `${id}-EndpointConfig`, {
       productionVariants: this.config.map((config, i) => ({

lib/osml/model_endpoint/me_test_endpoints.ts

@@ -73,6 +73,12 @@ export class METestEndpointsConfig extends BaseConfig {
    */
   public DEPLOY_SM_FLOOD_ENDPOINT: boolean;
 
+  /**
+   * Whether to deploy the SageMaker multi-container model endpoint.
+   * @default true
+   */
+  public DEPLOY_MULTI_CONTAINER_ENDPOINT: boolean;
+
   /**
    * The CPU allocation for the HTTP endpoint.
    * @default 4096
@@ -138,6 +144,12 @@ export class METestEndpointsConfig extends BaseConfig {
    */
   public SM_CENTER_POINT_MODEL: string;
 
+  /**
+   * The name of the multi-container SageMaker endpoint.
+   * @default "multi-container"
+   */
+  public SM_MULTI_CONTAINER_ENDPOINT: string;
+
   /**
    * The SageMaker CPU instance type.
    * @default "ml.m5.xlarge"
@@ -176,6 +188,7 @@ export class METestEndpointsConfig extends BaseConfig {
       DEPLOY_SM_AIRCRAFT_ENDPOINT: true,
       DEPLOY_SM_CENTERPOINT_ENDPOINT: true,
       DEPLOY_SM_FLOOD_ENDPOINT: true,
+      DEPLOY_MULTI_CONTAINER_ENDPOINT: true,
       HTTP_ENDPOINT_CPU: 4096,
       HTTP_ENDPOINT_CONTAINER_PORT: 8080,
       HTTP_ENDPOINT_DOMAIN_NAME: "test-http-model-endpoint",
@@ -185,8 +198,9 @@ export class METestEndpointsConfig extends BaseConfig {
       HTTP_ENDPOINT_MEMORY: 16384,
       SM_AIRCRAFT_MODEL: "aircraft",
       SM_CENTER_POINT_MODEL: "centerpoint",
-      SM_CPU_INSTANCE_TYPE: "ml.m5.xlarge",
       SM_FLOOD_MODEL: "flood",
+      SM_MULTI_CONTAINER_ENDPOINT: "multi-container",
+      SM_CPU_INSTANCE_TYPE: "ml.m5.xlarge",
       ...config
     });
   }
@@ -271,6 +285,11 @@ export class METestEndpoints extends Construct {
    */
   public aircraftModelEndpoint?: MESMEndpoint;
 
+  /**
+   * SM endpoint for testing a multi-container configuration.
+   */
+  public multiContainerModelEndpoint?: MESMEndpoint;
+
   /**
    * Security Group ID associated with the endpoints.
    */
@@ -474,5 +493,45 @@ export class METestEndpoints extends Construct {
       );
       this.aircraftModelEndpoint.node.addDependency(this.modelContainer);
     }
+
+    // Build a multi-container endpoint
+    if (this.config.DEPLOY_MULTI_CONTAINER_ENDPOINT) {
+      this.multiContainerModelEndpoint = new MESMEndpoint(
+        this,
+        "OSMLMultiContainerModelEndpoint",
+        {
+          containerImageUri: this.modelContainer.containerUri,
+          modelName: this.config.SM_MULTI_CONTAINER_ENDPOINT,
+          roleArn: this.smRole.roleArn,
+          instanceType: this.config.SM_CPU_INSTANCE_TYPE,
+          subnetIds: props.osmlVpc.selectedSubnets.subnetIds,
+          config: [
+            new MESMEndpointConfig({
+              SECURITY_GROUP_ID: this.securityGroupId,
+              CONTAINERS: [
+                {
+                  imageUri: this.modelContainer.containerUri,
+                  environment: {
+                    MODEL_SELECTION: this.config.SM_CENTER_POINT_MODEL
+                  },
+                  repositoryAccessMode:
+                    this.modelContainer.repositoryAccessMode,
+                  containerHostname: "centerpoint-container"
+                },
+                {
+                  imageUri: this.modelContainer.containerUri,
+                  environment: {
+                    MODEL_SELECTION: this.config.SM_AIRCRAFT_MODEL
+                  },
+                  repositoryAccessMode:
+                    this.modelContainer.repositoryAccessMode,
+                  containerHostname: "aircraft-container"
+                }
+              ]
+            })
+          ]
+        }
+      );
+    }
   }
 }