feat: enable support for gpu configuration (#50)

feat: enable support for gpu configuration


Refs: #40

---------

Signed-off-by: seven <[email protected]>

Author: Blankll

samples/aliyun-poc-fc-gpu.yml

@@ -28,6 +28,7 @@ functions:
       cmd: "npm start"
       port: 9000
     memory: 512
+    gpu: TESLA_8
     timeout: 10
     network:
       vpc_id: vpc-2vc8v9btc8470laqui9bk

src/parser/functionParser.ts

@@ -1,4 +1,4 @@
-import { FunctionDomain, FunctionRaw, NasStorageClassEnum } from '../types';
+import { FunctionDomain, FunctionGpuEnum, FunctionRaw, NasStorageClassEnum } from '../types';
 import { isEmpty } from 'lodash';
 
 export const parseFunction = (functions?: {
@@ -13,6 +13,7 @@ export const parseFunction = (functions?: {
     code: func.code,
     container: func.container,
     memory: func.memory,
+    gpu: func.gpu as FunctionGpuEnum,
     timeout: func.timeout,
     environment: func.environment,
     log: func.log,

src/stack/rosStack/function.ts

@@ -1,4 +1,10 @@
-import { ActionContext, FunctionDomain, NasStorageClassEnum, ServerlessIac } from '../../types';
+import {
+  ActionContext,
+  FunctionDomain,
+  FunctionGpuEnum,
+  NasStorageClassEnum,
+  ServerlessIac,
+} from '../../types';
 import {
   CODE_ZIP_SIZE_LIMIT,
   encodeBase64ForRosId,
@@ -25,13 +31,25 @@ const storageClassMap = {
   [NasStorageClassEnum.EXTREME_STANDARD]: { fileSystemType: 'extreme', storageType: 'standard' },
   [NasStorageClassEnum.EXTREME_ADVANCE]: { fileSystemType: 'extreme', storageType: 'advance' },
 };
+
 const securityGroupRangeMap: { [key: string]: string } = {
   TCP: '1/65535',
   UDP: '1/65535',
   ICMP: '-1/-1',
   GRE: '-1/-1',
   ALL: '-1/-1',
 };
+const gpuConfigMap = {
+  [FunctionGpuEnum.TESLA_8]: { gpuMemorySize: 8192, gpuType: 'fc.gpu.tesla.1' },
+  [FunctionGpuEnum.TESLA_12]: { gpuMemorySize: 12288, gpuType: 'fc.gpu.tesla.1' },
+  [FunctionGpuEnum.TESLA_16]: { gpuMemorySize: 16384, gpuType: 'fc.gpu.tesla.1' },
+  [FunctionGpuEnum.AMPERE_8]: { gpuMemorySize: 8192, gpuType: 'fc.gpu.ampere.1' },
+  [FunctionGpuEnum.AMPERE_12]: { gpuMemorySize: 12288, gpuType: 'fc.gpu.ampere.1' },
+  [FunctionGpuEnum.AMPERE_16]: { gpuMemorySize: 16384, gpuType: 'fc.gpu.ampere.1' },
+  [FunctionGpuEnum.AMPERE_24]: { gpuMemorySize: 24576, gpuType: 'fc.gpu.ampere.1' },
+  [FunctionGpuEnum.ADA_48]: { gpuMemorySize: 49152, gpuType: 'fc.gpu.ada.1' },
+};
+
 const transformSecurityRules = (rules: Array<string>, ruleType: 'INGRESS' | 'EGRESS') => {
   return rules.map((rule) => {
     const [protocol, cidrIp, portRange] = rule.split(':');
@@ -49,6 +67,13 @@ const transformSecurityRules = (rules: Array<string>, ruleType: 'INGRESS' | 'EGR
   });
 };
 
+const transformGpuConfig = (gpu: FunctionDomain['gpu']) => {
+  if (!gpu) {
+    return undefined;
+  }
+
+  return gpuConfigMap[gpu];
+};
 export const resolveFunctions = (
   scope: ros.Construct,
   functions: Array<FunctionDomain> | undefined,
@@ -248,8 +273,9 @@ export const resolveFunctions = (
       {
         functionName: replaceReference(fnc.name, context),
         memorySize: replaceReference(fnc.memory, context),
-        timeout: replaceReference(fnc.timeout, context),
         diskSize: fnc.storage?.disk,
+        gpuConfig: transformGpuConfig(fnc.gpu),
+        timeout: replaceReference(fnc.timeout, context),
         environmentVariables: replaceReference(fnc.environment, context),
         logConfig,
         vpcConfig,

src/types/domains/function.ts

@@ -11,6 +11,7 @@ export type FunctionRaw = {
     port: number;
   };
   memory: number;
+  gpu: string;
   timeout: number;
   log?: boolean;
   environment?: {
@@ -48,6 +49,7 @@ export type FunctionDomain = {
     port: number;
   };
   memory: number;
+  gpu?: FunctionGpuEnum;
   timeout: number;
   log?: boolean;
   environment?: {
@@ -77,3 +79,14 @@ export enum NasStorageClassEnum {
   EXTREME_STANDARD = 'EXTREME_STANDARD',
   EXTREME_ADVANCE = 'EXTREME_ADVANCE',
 }
+
+export enum FunctionGpuEnum {
+  TESLA_8 = 'TESLA_8',
+  TESLA_12 = 'TESLA_12',
+  TESLA_16 = 'TESLA_16',
+  AMPERE_8 = 'AMPERE_8',
+  AMPERE_12 = 'AMPERE_12',
+  AMPERE_16 = 'AMPERE_16',
+  AMPERE_24 = 'AMPERE_24',
+  ADA_48 = 'ADA_48',
+}

src/validator/functionSchema.ts

@@ -46,6 +46,19 @@ export const functionSchema = {
           },
         },
         memory: { type: 'number' },
+        gpu: {
+          type: 'string',
+          enum: [
+            'TESLA_8',
+            'TESLA_12',
+            'TESLA_16',
+            'AMPERE_8',
+            'AMPERE_12',
+            'AMPERE_16',
+            'AMPERE_24',
+            'ADA_48',
+          ],
+        },
         timeout: { type: 'number' },
         log: { type: 'boolean' },
         environment: {

tests/fixtures/deployFixture.ts

@@ -994,6 +994,32 @@ export const oneFcWithContainerRos = {
   },
 };
 
+export const oneFcWithGpuIac = {
+  ...oneFcIac,
+  functions: [
+    {
+      ...((oneFcIac.functions && oneFcIac.functions[0]) ?? {}),
+      gpu: 'TESLA_8',
+    },
+  ],
+} as ServerlessIac;
+export const oneFcWithGpuRos = {
+  ...oneFcRos,
+  Resources: {
+    ...oneFcRos.Resources,
+    hello_fn: {
+      ...oneFcRos.Resources.hello_fn,
+      Properties: {
+        ...oneFcRos.Resources.hello_fn.Properties,
+        GpuConfig: {
+          GpuMemorySize: 8192,
+          GpuType: 'fc.gpu.tesla.1',
+        },
+      },
+    },
+  },
+};
+
 export const largeCodeRos = {
   Description: 'my-demo-service stack',
   Mappings: {

tests/stack/deploy.test.ts

@@ -19,6 +19,8 @@ import {
   oneFcRos,
   oneFcWithContainerIac,
   oneFcWithContainerRos,
+  oneFcWithGpuIac,
+  oneFcWithGpuRos,
   oneFcWithStageRos,
   referredServiceIac,
   referredServiceRos,
@@ -217,7 +219,7 @@ describe('Unit tests for stack deployment', () => {
     });
   });
 
-  describe('unit test for function nas attachment', () => {
+  describe('unit test for serverless Gpu', () => {
     it('should deploy function with nas when nas field is provided', async () => {
       const stackName = 'my-demo-stack-with-nas';
       mockedRosStackDeploy.mockResolvedValue(stackName);
@@ -231,9 +233,7 @@ describe('Unit tests for stack deployment', () => {
         { stackName },
       ]);
     });
-  });
 
-  describe('unit test for function with container', () => {
     it('should deploy function with container when container field is provided', async () => {
       const stackName = 'my-demo-stack-with-container';
       mockedRosStackDeploy.mockResolvedValue(stackName);
@@ -247,5 +247,19 @@ describe('Unit tests for stack deployment', () => {
         { stackName },
       ]);
     });
+
+    it('should deploy function with gpu configured', async () => {
+      const stackName = 'my-demo-stack-with-gpu';
+      mockedRosStackDeploy.mockResolvedValue(stackName);
+
+      await deployStack(stackName, oneFcWithGpuIac, { stackName } as ActionContext);
+
+      expect(mockedRosStackDeploy).toHaveBeenCalledTimes(2);
+      expect(mockedRosStackDeploy.mock.calls[1]).toEqual([
+        stackName,
+        oneFcWithGpuRos,
+        { stackName },
+      ]);
+    });
   });
 });