feat: add gpus_per_replicas

src/pages/llmodels/config/index.ts

@@ -194,6 +194,11 @@ export const ScheduleValueMap = {
   SpecificGPUType: 'specific_gpu_type'
 };
 
+export const gpusCountTypeMap = {
+  Auto: 'auto',
+  Custom: 'custom'
+};
+
 export const scheduleList = [
   {
     label: 'models.form.scheduletype.auto',
@@ -324,6 +329,7 @@ export const excludeFields = [
   'backend_version',
   'ollama_library_model_name',
   'scheduleType',
+  'gpusCountType',
   'placement_strategy',
   'backend',
   'gpu_selector',

src/pages/llmodels/config/types.ts

@@ -27,6 +27,7 @@ export interface ListItem {
   access_policy: 'public' | 'authed' | 'allowed_users';
   gpu_selector?: {
     gpu_ids: string[];
+    gpus_per_replica?: number;
   };
   worker_selector?: object;
 }
@@ -63,11 +64,13 @@ export interface FormData {
     gpu_ids?: string[];
     gpu_type?: string;
     gpu_count?: number;
+    gpus_per_replica?: number;
   };
   placement_strategy?: string;
   cpu_offloading?: boolean;
   worker_selector?: object;
   scheduleType?: string;
+  gpusCountType?: string;
   name: string;
   replicas: number;
   description: string;
@@ -191,6 +194,7 @@ export interface CatalogSpec {
   worker_selector: Record<string, any>;
   gpu_selector: {
     gpu_ids: string[];
+    gpus_per_replica: number;
   };
   backend: string;
   backend_version: string;
@@ -220,6 +224,7 @@ export interface EvaluateSpec {
   worker_selector?: Record<string, any>;
   gpu_selector?: {
     gpu_ids: string[];
+    gpus_per_replica: number;
   };
   backend?: string;
   backend_version?: string;

src/pages/llmodels/config/utils.ts

@@ -1,4 +1,5 @@
 import _ from 'lodash';
+import { gpusCountTypeMap } from '.';
 import { backendOptionsMap } from '../config/backend-parameters';
 import { FormData } from './types';
 
@@ -33,6 +34,11 @@ export const generateGPUSelector = (data: any, gpuOptions: any[]) => {
   };
 };
 
+/**
+ * before submit the form, generate the gpu_selector field
+ * @param data
+ * @returns
+ */
 export const generateGPUIds = (data: FormData) => {
   const gpu_ids = _.get(data, 'gpu_selector.gpu_ids', []);
   console.log('generateGPUIds', gpu_ids);
@@ -54,10 +60,17 @@ export const generateGPUIds = (data: FormData) => {
     },
     []
   );
-
+  if (gpusCountTypeMap.Auto === data.gpusCountType) {
+    return {
+      gpu_selector: {
+        gpu_ids: result || []
+      }
+    };
+  }
   return {
     gpu_selector: {
-      gpu_ids: result || []
+      gpu_ids: result || [],
+      gpus_per_replica: data.gpu_selector?.gpus_per_replica || null
     }
   };
 };

src/pages/llmodels/forms/index.tsx

@@ -5,7 +5,7 @@ import { useIntl } from '@umijs/max';
 import { Form } from 'antd';
 import _ from 'lodash';
 import React, { forwardRef, useImperativeHandle } from 'react';
-import { excludeFields, ScheduleValueMap } from '../config';
+import { excludeFields, gpusCountTypeMap, ScheduleValueMap } from '../config';
 import { backendOptionsMap } from '../config/backend-parameters';
 import { FormContext } from '../config/form-context';
 import {
@@ -120,7 +120,7 @@ const DataForm: React.FC<DataFormProps> = forwardRef((props, ref) => {
     data.categories = data.categories ? [data.categories] : [];
     const gpuSelector = generateGPUIds(data);
     const allValues = {
-      ..._.omit(data, ['scheduleType']),
+      ..._.omit(data, ['scheduleType', 'gpusCountType']),
       ...gpuSelector
     };
 
@@ -250,6 +250,7 @@ const DataForm: React.FC<DataFormProps> = forwardRef((props, ref) => {
           placement_strategy: 'spread',
           cpu_offloading: true,
           scheduleType: ScheduleValueMap.Auto,
+          gpusCountType: gpusCountTypeMap.Auto,
           categories: null,
           restart_on_error: true,
           distributed_inference_across_workers: true,

src/pages/llmodels/forms/schedule-type.tsx

@@ -1,3 +1,4 @@
+import SealInputNumber from '@/components/seal-form/input-number';
 import SealCascader from '@/components/seal-form/seal-cascader';
 import SealSelect from '@/components/seal-form/seal-select';
 import TooltipList from '@/components/tooltip-list';
@@ -6,9 +7,9 @@ import { useIntl } from '@umijs/max';
 import { Form } from 'antd';
 import React from 'react';
 import GPUCard from '../components/gpu-card';
-import { scheduleList, ScheduleValueMap } from '../config';
+import { gpusCountTypeMap, scheduleList, ScheduleValueMap } from '../config';
 import { backendOptionsMap } from '../config/backend-parameters';
-import { useCatalogFormContext, useFormContext } from '../config/form-context';
+import { useFormContext } from '../config/form-context';
 
 const scheduleTypeTips = [
   {
@@ -27,13 +28,31 @@ const scheduleTypeTips = [
   }
 ];
 
+const gpuAllocateTypeTips = [
+  {
+    title: {
+      text: 'models.form.gpusAllocationType.auto',
+      locale: true
+    },
+    tips: 'models.form.gpusAllocationType.auto.tips'
+  },
+  {
+    title: {
+      text: 'models.form.gpusAllocationType.custom',
+      locale: true
+    },
+    tips: 'models.form.gpusAllocationType.custom.tips'
+  }
+];
+
 const ScheduleTypeForm: React.FC = () => {
   const intl = useIntl();
   const { onValuesChange, gpuOptions } = useFormContext();
-  const { onQuantizationChange } = useCatalogFormContext();
   const { getRuleMessage } = useAppUtils();
   const form = Form.useFormInstance();
   const scheduleType = Form.useWatch('scheduleType', form);
+  const gpusCountType = Form.useWatch('gpusCountType', form);
+  const gpuSelectorIds = Form.useWatch(['gpu_selector', 'gpu_ids'], form);
 
   const handleScheduleTypeChange = (value: string) => {
     if (value === ScheduleValueMap.Auto) {
@@ -41,8 +60,11 @@ const ScheduleTypeForm: React.FC = () => {
     }
   };
 
-  const handleOnQuantizationChange = (val: any) => {
-    onQuantizationChange?.(val);
+  const handleGpusCountTypeChange = (val: string) => {
+    if (val === 'custom') {
+      form.setFieldValue(['gpu_selector', 'gpus_per_replica'], 2);
+    }
+    onValuesChange?.({}, form.getFieldsValue());
   };
 
   const handleBeforeGpuSelectorChange = (gpuIds: any[]) => {};
@@ -52,6 +74,23 @@ const ScheduleTypeForm: React.FC = () => {
     onValuesChange?.({}, form.getFieldsValue());
   };
 
+  const handleOnStepReplicaStep = (
+    value: number,
+    info: { offset: number; type: 'up' | 'down' }
+  ) => {
+    let newValue = value;
+    const isPowerOfTwo = (n: number) => (n & (n - 1)) === 0 && n !== 0; // check power of two
+    if (!isPowerOfTwo(value)) {
+      if (info.type === 'up') {
+        newValue = Math.pow(2, Math.ceil(Math.log2(value)));
+      } else {
+        newValue = Math.pow(2, Math.floor(Math.log2(value)));
+      }
+    }
+    form.setFieldValue(['gpu_selector', 'gpus_per_replica'], newValue);
+    onValuesChange?.({}, form.getFieldsValue());
+  };
+
   return (
     <>
       <Form.Item name="scheduleType">
@@ -115,6 +154,43 @@ const ScheduleTypeForm: React.FC = () => {
                 onChange={handleGpuSelectorChange}
               ></SealCascader>
             </Form.Item>
+            <Form.Item name="gpusCountType">
+              <SealSelect
+                onChange={handleGpusCountTypeChange}
+                label={intl.formatMessage({
+                  id: 'models.form.gpusAllocationType'
+                })}
+                description={
+                  <TooltipList list={gpuAllocateTypeTips}></TooltipList>
+                }
+                options={[
+                  {
+                    label: intl.formatMessage({
+                      id: 'models.form.gpusAllocationType.auto'
+                    }),
+                    value: gpusCountTypeMap.Auto
+                  },
+                  {
+                    label: intl.formatMessage({
+                      id: 'models.form.gpusAllocationType.custom'
+                    }),
+                    value: gpusCountTypeMap.Custom
+                  }
+                ]}
+              ></SealSelect>
+            </Form.Item>
+            {gpusCountType === gpusCountTypeMap.Custom && (
+              <Form.Item name={['gpu_selector', 'gpus_per_replica']}>
+                <SealInputNumber
+                  label={intl.formatMessage({
+                    id: 'models.form.gpusperreplica'
+                  })}
+                  min={1}
+                  step={1}
+                  onStep={handleOnStepReplicaStep}
+                />
+              </Form.Item>
+            )}
           </>
         )}
     </>

src/pages/llmodels/hooks/use-form-initial-values.ts

@@ -10,7 +10,7 @@ import { ListItem as WorkerListItem } from '@/pages/resources/config/types';
 import { useAtom } from 'jotai';
 import { useState } from 'react';
 import { queryGPUList } from '../apis';
-import { ScheduleValueMap } from '../config';
+import { gpusCountTypeMap, ScheduleValueMap } from '../config';
 import { GPUListItem, ListItem } from '../config/types';
 
 type EmptyObject = Record<never, never>;
@@ -226,13 +226,22 @@ export default function useFormInitialValues() {
     }
   };
 
+  /**
+   * before set the form initial values, generate the form values
+   * @param data
+   * @param gpuOptions
+   * @returns
+   */
   const generateFormValues = (data: ListItem, gpuOptions: any[]) => {
     const formData = {
       ...data,
       categories: data?.categories?.length ? data.categories[0] : null,
       scheduleType: data?.gpu_selector
         ? ScheduleValueMap.Manual
-        : ScheduleValueMap.Auto
+        : ScheduleValueMap.Auto,
+      gpusCountType: data?.gpu_selector?.gpus_per_replica
+        ? gpusCountTypeMap.Custom
+        : gpusCountTypeMap.Auto
     };
     return formData;
   };