feat(FR-1448): allow user to select multi agents when creating multi-node session

Author: nowgnuesLee

react/src/components/AgentSelect.tsx

@@ -9,7 +9,7 @@ import React, { useDeferredValue, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { graphql, useLazyLoadQuery } from 'react-relay';
 
-interface Props extends SelectProps {
+interface Props extends Omit<SelectProps, 'options'> {
   autoSelectDefault?: boolean;
   fetchKey?: string;
   resourceGroup?: string | null;
@@ -122,19 +122,33 @@ const AgentSelect: React.FC<Props> = ({
     : undefined;
   return (
     <Select
-      onChange={(value, option) => {
-        setValue(value, option);
-      }}
       loading={searchStr !== deferredSearchStr}
       filterOption={false}
       showSearch
       searchValue={searchStr}
       onSearch={(v) => {
         setSearchStr(v);
       }}
-      {...selectProps}
-      value={value}
       options={filterOutEmpty([autoSelectIfMatch, ...agentOptions])}
+      //override props.onChange and props.value, it is handled by useControllableValue
+      {..._.omit(selectProps, ['value', 'onChange'])}
+      onChange={(value: unknown, option) => {
+        if (
+          selectProps.mode === 'multiple' &&
+          _.isArray(value) &&
+          _.isArray(option)
+        ) {
+          if (_.last(value) === 'auto' || value.length === 0) {
+            value = ['auto'];
+            option = _.last(option);
+          } else if (value[0] === 'auto' && value.length > 1) {
+            value = value.slice(1);
+            option = option.slice(1);
+          }
+        }
+        setValue(value, option);
+      }}
+      value={value}
     />
   );
 };

react/src/components/SessionFormItems/ResourceAllocationFormItems.tsx

@@ -46,7 +46,7 @@ export const RESOURCE_ALLOCATION_INITIAL_FORM_VALUES: DeepPartial<ResourceAlloca
     cluster_mode: 'single-node',
     cluster_size: 1,
     enabledAutomaticShmem: true,
-    agent: 'auto',
+    agent: ['auto'],
   };
 
 export const isMinOversMaxValue = (min: number, max: number) => {
@@ -67,7 +67,7 @@ export interface ResourceAllocationFormValue {
   cluster_size: number;
   enabledAutomaticShmem: boolean;
   allocationPreset?: string;
-  agent?: string;
+  agent?: string[] | string;
 }
 
 export type MergedResourceAllocationFormValue = ResourceAllocationFormValue &
@@ -1233,14 +1233,26 @@ const ResourceAllocationFormItems: React.FC<
                 <AgentSelect
                   resourceGroup={currentResourceGroupInForm}
                   fetchKey={agentFetchKey}
+                  mode={
+                    baiClient.supports('multi-agents') ? 'multiple' : undefined
+                  }
+                  labelRender={
+                    baiClient.supports('multi-agents')
+                      ? ({ label, value }) => {
+                          return value === 'auto' ? label : value;
+                        }
+                      : undefined
+                  }
                   onChange={(value) => {
-                    if (value !== 'auto') {
+                    if (
+                      !baiClient.supports('multi-agents') &&
+                      (!_.isEqual(value, 'auto') || !_.isEqual(value, ['auto']))
+                    ) {
                       form.setFieldsValue({
                         cluster_mode: 'single-node',
                         cluster_size: 1,
                       });
                     }
-                    // TODO: set cluster mode to single node and cluster size to 1 when agent value is not "auto"
                   }}
                 ></AgentSelect>
               </Form.Item>
@@ -1286,7 +1298,11 @@ const ResourceAllocationFormItems: React.FC<
                       onChange={() => {
                         form.validateFields().catch(() => {});
                       }}
-                      disabled={getFieldValue('agent') !== 'auto'}
+                      disabled={
+                        !baiClient.supports('multi-agents') &&
+                        !_.isEqual(getFieldValue('agent'), 'auto') &&
+                        !_.isEqual(getFieldValue('agent'), ['auto'])
+                      }
                     >
                       <Radio.Button value="single-node">
                         {t('session.launcher.SingleNode')}
@@ -1351,7 +1367,9 @@ const ResourceAllocationFormItems: React.FC<
                             }
                             disabled={
                               derivedClusterSizeMaxLimit === 1 ||
-                              getFieldValue('agent') !== 'auto'
+                              (!baiClient.supports('multi-agents') &&
+                                !_.isEqual(getFieldValue('agent'), 'auto') &&
+                                !_.isEqual(getFieldValue('agent'), ['auto']))
                             }
                             sliderProps={{
                               marks: {

react/src/hooks/useStartSession.tsx

@@ -181,16 +181,33 @@ export const useStartSession = () => {
           preopen_ports: transformPortValuesToNumbers(values.ports),
 
           // Agent selection (optional)
-          ...(baiClient.supports('agent-select') &&
-          !baiClient?._config?.hideAgents &&
-          values.agent !== 'auto'
-            ? {
-                // Filter out undefined values
-                agent_list: [values.agent].filter(
-                  (agent): agent is string => !!agent,
+          ...(() => {
+            if (values.agent === undefined) return {};
+            const agents = _.castArray(values.agent);
+            if (
+              !baiClient.supports('agent-select') ||
+              baiClient?._config?.hideAgents ||
+              _.isEqual(agents, ['auto']) ||
+              agents.length === 0
+            ) {
+              return {};
+            }
+            if (
+              !baiClient.supports('multi-agents') &&
+              values.cluster_mode === 'multi-node'
+            ) {
+              // The server now requires agents equivalent to the cluster size.
+              return {
+                agent_list: Array.from(
+                  { length: values.cluster_size },
+                  () => agents[0],
                 ),
-              }
-            : undefined),
+              };
+            }
+            return {
+              agent_list: _.filter(agents, (item) => !!item),
+            };
+          })(),
         },
       },
     };

react/src/pages/SessionLauncherPage.tsx

@@ -1281,7 +1281,7 @@ const SessionLauncherPage = () => {
                   command: undefined,
                   scheduleDate: undefined,
                 },
-                agent: 'auto', // Add the missing 'agent' property
+                agent: ['auto'], // Add the missing 'agent' property
               } as SessionLauncherFormData,
               formValue,
             );

src/lib/backend.ai-client-esm.ts

@@ -825,6 +825,9 @@ class Client {
     if (this.isManagerVersionCompatibleWith('25.15.0')) {
       this._features['agent-stats'] = true;
     }
+    if (this.isManagerVersionCompatibleWith('25.16.0')) {
+      this._features['multi-agents'] = true;
+    }
   }
 
   /**