Skip to content

Commit

Permalink
sg override fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
@bedanley
bedanley authored Nov 26, 2024
1 parent d39a05f commit a9e6453
Show file tree
Hide file tree
Showing 5 changed files with 42 additions and 28 deletions.
24 changes: 18 additions & 6 deletions lib/core/iam/SecurityGroups.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,25 @@
*/

/**
* List of all security groups used for overrides
* List of all security group ids used for overrides
*/
export enum SecurityGroups {
LITE_LLM_SG = 'LiteLLMScalingSg',
export enum SecurityGroupEnum {
LITE_LLM_SG = 'LISA-LiteLLMScalingSg',
ECS_MODEL_ALB_SG = 'EcsModelAlbSg',
REST_API_ALB_SG = 'RestApiAlbSg',
LAMBDA_SG = 'LambdaSg',
OPEN_SEARCH_SG = 'OpenSearchSg',
PG_VECTOR_SG = 'PGVectorSg',
LAMBDA_SG = 'LambdaSecurityGroup',
OPEN_SEARCH_SG = 'LISA-OpenSearchSg',
PG_VECTOR_SG = 'LISA-PGVectorSg',
}

/**
* List of all security group names used for overrides.
* LiteLLMScalingSg does not have a predefined name
*/
export const SecurityGroupNames: Record<string, string> = {
'EcsModelAlbSg' : 'ECS-ALB-SG',
'RestApiAlbSg' : 'RestAPI-ALB-SG',
'LambdaSecurityGroup' : 'Lambda-SG',
'LISA-OpenSearchSg' : 'OpenSearch-SG',
'LISA-PGVectorSg' : 'LISA-PGVector-SG',
};
10 changes: 5 additions & 5 deletions lib/networking/vpc/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ import { Construct } from 'constructs';

import { createCdkId } from '../../core/utils';
import { BaseProps, SecurityGroups } from '../../schema';
import { SecurityGroups as SecurityGroupsEnum } from '../../core/iam/SecurityGroups';
import { SecurityGroupEnum } from '../../core/iam/SecurityGroups';
import { SubnetGroup } from 'aws-cdk-lib/aws-rds';
import { SecurityGroupFactory } from './security-group-factory';

Expand Down Expand Up @@ -127,7 +127,7 @@ export class Vpc extends Construct {
const ecsModelAlbSg = SecurityGroupFactory.createSecurityGroup(
this,
sgOverrides?.modelSecurityGroupId,
SecurityGroupsEnum.ECS_MODEL_ALB_SG,
SecurityGroupEnum.ECS_MODEL_ALB_SG,
config.deploymentName,
vpc,
'ECS model application load balancer',
Expand All @@ -139,7 +139,7 @@ export class Vpc extends Construct {
const restApiAlbSg = SecurityGroupFactory.createSecurityGroup(
this,
sgOverrides?.restAlbSecurityGroupId,
SecurityGroupsEnum.REST_API_ALB_SG,
SecurityGroupEnum.REST_API_ALB_SG,
config.deploymentName,
vpc,
'REST API application load balancer',
Expand All @@ -155,7 +155,7 @@ export class Vpc extends Construct {
const lambdaSg = SecurityGroupFactory.createSecurityGroup(
this,
sgOverrides?.lambdaSecurityGroupId,
SecurityGroupsEnum.LAMBDA_SG,
SecurityGroupEnum.LAMBDA_SG,
config.deploymentName,
vpc,
'authorizer and API Lambdas',
Expand All @@ -176,6 +176,6 @@ export class Vpc extends Construct {
new CfnOutput(this, 'vpcCidrBlock', { value: vpc.vpcCidrBlock });
new CfnOutput(this, 'ecsModelAlbSg', { value: ecsModelAlbSg.securityGroupId });
new CfnOutput(this, 'restApiAlbSg', { value: restApiAlbSg.securityGroupId });
new CfnOutput(this, 'lambdaSg', { value: lambdaSg.securityGroupId });
new CfnOutput(this, 'lambdaSecurityGroup', { value: lambdaSg.securityGroupId });
}
}
16 changes: 9 additions & 7 deletions lib/networking/vpc/security-group-factory.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import { ISecurityGroup, IVpc, Peer, Port, SecurityGroup } from 'aws-cdk-lib/aws-ec2';
import { Config } from '../../schema';
import { createCdkId } from '../../core/utils';
import { SecurityGroupNames } from '../../core/iam/SecurityGroups';
import { Vpc } from '.';
import { IConstruct } from 'constructs';

Expand All @@ -29,7 +30,7 @@ export class SecurityGroupFactory {
* Creates a security group for the VPC.
*
* @param securityGroupOverride - security group override
* @param {string} securityGroupName - The name of the security group.
* @param {string} securityGroupId - The name of the security group.
* @param {string} deploymentName - The deployment name.
* @param {Vpc} vpc - The virtual private cloud.
* @param {string} description - The description of the security group.
Expand All @@ -38,24 +39,25 @@ export class SecurityGroupFactory {
static createSecurityGroup (
construct: IConstruct,
securityGroupOverride: string | undefined,
securityGroupName: string,
deploymentName: string,
securityGroupId: string,
deploymentName: string | undefined,
vpc: IVpc,
description: string,
): ISecurityGroup {
if (securityGroupOverride) {
console.log(`Security Role Override provided. Using ${securityGroupOverride} for ${securityGroupName}`);
const sg = SecurityGroup.fromSecurityGroupId(construct, securityGroupName, securityGroupOverride);
console.log(`Security Role Override provided. Using ${securityGroupOverride} for ${securityGroupId}`);
const sg = SecurityGroup.fromSecurityGroupId(construct, securityGroupId, securityGroupOverride);
// Validate the security group exists
if (!sg) {
throw new Error(`Security group ${sg} not found`);
}
return sg;
} else {
const sg = new SecurityGroup(construct, securityGroupName, {
securityGroupName: createCdkId([deploymentName, securityGroupName]),
const securityGroupName = SecurityGroupNames[securityGroupId];
const sg = new SecurityGroup(construct, securityGroupId, {
vpc: vpc,
description: `Security group for ${description}`,
...(securityGroupName && {securityGroupName: createCdkId(deploymentName ? [deploymentName, securityGroupName] : [securityGroupName])}),
});

return sg;
Expand Down
12 changes: 6 additions & 6 deletions lib/rag/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ import { Layer } from '../core/layers';
import { createCdkId } from '../core/utils';
import { Vpc } from '../networking/vpc';
import { BaseProps, RagRepositoryType } from '../schema';
import { SecurityGroups } from '../core/iam/SecurityGroups';
import { SecurityGroupEnum } from '../core/iam/SecurityGroups';
import { SecurityGroupFactory } from '../networking/vpc/security-group-factory';

import { IngestPipelineStateMachine } from './state_machine/ingest-pipeline';
Expand Down Expand Up @@ -154,13 +154,13 @@ export class LisaRagStack extends Stack {
const openSearchSg = SecurityGroupFactory.createSecurityGroup(
this,
config.securityGroupConfig?.openSearchSecurityGroupId,
SecurityGroups.OPEN_SEARCH_SG,
SecurityGroupEnum.OPEN_SEARCH_SG,
config.deploymentName,
vpc.vpc,
'RAG OpenSearch domain',
);
if (!config.securityGroupConfig?.openSearchSecurityGroupId) {
SecurityGroupFactory.addIngress(openSearchSg, SecurityGroups.OPEN_SEARCH_SG, vpc, config);
SecurityGroupFactory.addIngress(openSearchSg, SecurityGroupEnum.OPEN_SEARCH_SG, vpc, config);
}

registeredRepositories.push({ repositoryId: ragConfig.repositoryId, type: ragConfig.type });
Expand Down Expand Up @@ -273,13 +273,13 @@ export class LisaRagStack extends Stack {
const pgvectorSg = SecurityGroupFactory.createSecurityGroup(
this,
config.securityGroupConfig?.pgVectorSecurityGroupId,
SecurityGroups.PG_VECTOR_SG,
config.deploymentName,
SecurityGroupEnum.PG_VECTOR_SG,
undefined,
vpc.vpc,
'RAG PGVector database',
);
if (!config.securityGroupConfig?.pgVectorSecurityGroupId) {
SecurityGroupFactory.addIngress(pgvectorSg, SecurityGroups.PG_VECTOR_SG, vpc, config);
SecurityGroupFactory.addIngress(pgvectorSg, SecurityGroupEnum.PG_VECTOR_SG, vpc, config);
}

const username = ragConfig.rdsConfig.username;
Expand Down
8 changes: 4 additions & 4 deletions lib/serve/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ import { Vpc } from '../networking/vpc';
import { BaseProps } from '../schema';
import { Effect, Policy, PolicyStatement } from 'aws-cdk-lib/aws-iam';
import { Secret } from 'aws-cdk-lib/aws-secretsmanager';
import { SecurityGroups } from '../core/iam/SecurityGroups';
import { SecurityGroupEnum } from '../core/iam/SecurityGroups';
import { SecurityGroupFactory } from '../networking/vpc/security-group-factory';

const HERE = path.resolve(__dirname);
Expand Down Expand Up @@ -147,13 +147,13 @@ export class LisaServeApplicationStack extends Stack {
const litellmDbSg = SecurityGroupFactory.createSecurityGroup(
this,
config.securityGroupConfig?.liteLlmDbSecurityGroupId,
SecurityGroups.LITE_LLM_SG,
config.deploymentName,
SecurityGroupEnum.LITE_LLM_SG,
undefined,
vpc.vpc,
'LiteLLM dynamic model management database',
);
if (!config.securityGroupConfig?.liteLlmDbSecurityGroupId) {
SecurityGroupFactory.addIngress(litellmDbSg, SecurityGroups.LITE_LLM_SG, vpc, config);
SecurityGroupFactory.addIngress(litellmDbSg, SecurityGroupEnum.LITE_LLM_SG, vpc, config);
}

const username = config.restApiConfig.rdsConfig.username;
Expand Down

0 comments on commit a9e6453

Please sign in to comment.