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(_.castArray(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,12 @@ const ResourceAllocationFormItems: React.FC<
                       onChange={() => {
                         form.validateFields().catch(() => {});
                       }}
-                      disabled={getFieldValue('agent') !== 'auto'}
+                      disabled={
+                        !baiClient.supports('multi-agents') &&
+                        !_.isEqual(_.castArray(getFieldValue('agent')), [
+                          'auto',
+                        ])
+                      }
                     >
                       <Radio.Button value="single-node">
                         {t('session.launcher.SingleNode')}
@@ -1351,7 +1368,11 @@ const ResourceAllocationFormItems: React.FC<
                             }
                             disabled={
                               derivedClusterSizeMaxLimit === 1 ||
-                              getFieldValue('agent') !== 'auto'
+                              (!baiClient.supports('multi-agents') &&
+                                !_.isEqual(
+                                  _.castArray(getFieldValue('agent')),
+                                  ['auto'],
+                                ))
                             }
                             sliderProps={{
                               marks: {

react/src/components/SessionLauncherPreview.tsx

@@ -523,7 +523,7 @@ const SessionLauncherPreview: React.FC<{
             {baiClient.supports('agent-select') &&
               !baiClient?._config?.hideAgents && (
                 <Descriptions.Item label={t('session.launcher.AgentNode')}>
-                  {form.getFieldValue('agent') ||
+                  {_.castArray(form.getFieldValue('agent')).join(', ') ||
                     t('session.launcher.AutoSelect')}
                 </Descriptions.Item>
               )}

react/src/hooks/useStartSession.tsx

@@ -251,12 +251,14 @@ export const useStartSession = () => {
           // Agent selection (optional)
           ...(baiClient.supports('agent-select') &&
           !baiClient?._config?.hideAgents &&
-          values.agent !== 'auto'
+          values.agent !== undefined &&
+          !_.isEqual(values.agent, ['auto'])
             ? {
                 // Filter out undefined values
-                agent_list: [values.agent].filter(
-                  (agent): agent is string => !!agent,
-                ),
+                agent_list: _.chain(values.agent)
+                  .castArray()
+                  .filter((agent): agent is string => !!agent)
+                  .value(),
               }
             : undefined),
         },

react/src/pages/SessionLauncherPage.tsx

@@ -1278,7 +1278,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;
+    }
   }
 
   /**