Skip to content

Commit a30d7fe

Browse files
nowgnuesLeeyomybaby
authored andcommitted
feat(FR-1448): allow user to select multi agents when creating multi-node session
1 parent 22982f8 commit a30d7fe

File tree

4 files changed

+59
-29
lines changed

4 files changed

+59
-29
lines changed

react/src/components/AgentSelect.tsx

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import React, { useDeferredValue, useState } from 'react';
99
import { useTranslation } from 'react-i18next';
1010
import { graphql, useLazyLoadQuery } from 'react-relay';
1111

12-
interface Props extends SelectProps {
12+
interface Props extends Omit<SelectProps, 'options'> {
1313
autoSelectDefault?: boolean;
1414
fetchKey?: string;
1515
resourceGroup?: string | null;
@@ -122,8 +122,17 @@ const AgentSelect: React.FC<Props> = ({
122122
: undefined;
123123
return (
124124
<Select
125-
onChange={(value, option) => {
125+
onChange={(value: unknown, option) => {
126+
if (Array.isArray(value)) {
127+
// multi-mode
128+
if (value[value.length - 1] === 'auto' || value.length === 0) {
129+
value = ['auto'];
130+
} else if (value[0] === 'auto' && value.length > 1) {
131+
value = value.slice(1);
132+
}
133+
}
126134
setValue(value, option);
135+
selectProps.onChange?.(value, option);
127136
}}
128137
loading={searchStr !== deferredSearchStr}
129138
filterOption={false}
@@ -132,7 +141,7 @@ const AgentSelect: React.FC<Props> = ({
132141
onSearch={(v) => {
133142
setSearchStr(v);
134143
}}
135-
{...selectProps}
144+
{..._.omit(selectProps, ['onChange'])}
136145
value={value}
137146
options={filterOutEmpty([autoSelectIfMatch, ...agentOptions])}
138147
/>

react/src/components/SessionFormItems/ResourceAllocationFormItems.tsx

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ export const RESOURCE_ALLOCATION_INITIAL_FORM_VALUES: DeepPartial<ResourceAlloca
4646
cluster_mode: 'single-node',
4747
cluster_size: 1,
4848
enabledAutomaticShmem: true,
49-
agent: 'auto',
49+
agent: ['auto'],
5050
};
5151

5252
export const isMinOversMaxValue = (min: number, max: number) => {
@@ -67,7 +67,7 @@ export interface ResourceAllocationFormValue {
6767
cluster_size: number;
6868
enabledAutomaticShmem: boolean;
6969
allocationPreset?: string;
70-
agent?: string;
70+
agent?: string[] | string;
7171
}
7272

7373
export type MergedResourceAllocationFormValue = ResourceAllocationFormValue &
@@ -1232,15 +1232,16 @@ const ResourceAllocationFormItems: React.FC<
12321232
<AgentSelect
12331233
resourceGroup={currentResourceGroupInForm}
12341234
fetchKey={agentFetchKey}
1235-
onChange={(value, option) => {
1236-
if (value !== 'auto') {
1237-
form.setFieldsValue({
1238-
cluster_mode: 'single-node',
1239-
cluster_size: 1,
1240-
});
1241-
}
1242-
// TODO: set cluster mode to single node and cluster size to 1 when agent value is not "auto"
1243-
}}
1235+
mode={
1236+
baiClient.supports('multi-agents') ? 'multiple' : undefined
1237+
}
1238+
labelRender={
1239+
baiClient.supports('multi-agents')
1240+
? ({ label, value }) => {
1241+
return value === 'auto' ? label : value;
1242+
}
1243+
: undefined
1244+
}
12441245
></AgentSelect>
12451246
</Form.Item>
12461247
</Suspense>
@@ -1285,7 +1286,6 @@ const ResourceAllocationFormItems: React.FC<
12851286
onChange={(e) => {
12861287
form.validateFields().catch(() => {});
12871288
}}
1288-
disabled={getFieldValue('agent') !== 'auto'}
12891289
>
12901290
<Radio.Button value="single-node">
12911291
{t('session.launcher.SingleNode')}
@@ -1348,10 +1348,7 @@ const ResourceAllocationFormItems: React.FC<
13481348
? derivedClusterSizeMaxLimit
13491349
: undefined
13501350
}
1351-
disabled={
1352-
derivedClusterSizeMaxLimit === 1 ||
1353-
getFieldValue('agent') !== 'auto'
1354-
}
1351+
disabled={derivedClusterSizeMaxLimit === 1}
13551352
sliderProps={{
13561353
marks: {
13571354
1: '1',

react/src/pages/SessionLauncherPage.tsx

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -549,16 +549,37 @@ const SessionLauncherPage = () => {
549549
preopen_ports: transformPortValuesToNumbers(values.ports),
550550

551551
// Agent selection (optional)
552-
...(baiClient.supports('agent-select') &&
553-
!baiClient?._config?.hideAgents &&
554-
values.agent !== 'auto'
555-
? {
556-
// Filter out undefined values
557-
agent_list: [values.agent].filter(
558-
(agent): agent is string => !!agent,
552+
...(() => {
553+
if (values.agent === undefined) return {};
554+
555+
const agents = _.castArray(values.agent);
556+
557+
if (
558+
!baiClient.supports('agent-select') ||
559+
baiClient?._config?.hideAgents ||
560+
_.isEqual(agents, ['auto']) ||
561+
agents.length === 0
562+
) {
563+
return {};
564+
}
565+
566+
if (
567+
!baiClient.supports('multi-agents') &&
568+
values.cluster_mode === 'multi-node'
569+
) {
570+
// The server now requires agents equivalent to the cluster size.
571+
return {
572+
agent_list: Array.from(
573+
{ length: values.cluster_size },
574+
() => agents[0],
559575
),
560-
}
561-
: undefined),
576+
};
577+
}
578+
579+
return {
580+
agent_list: _.filter(agents, (item) => !!item),
581+
};
582+
})(),
562583
},
563584
},
564585
};
@@ -1546,7 +1567,7 @@ const SessionLauncherPage = () => {
15461567
command: undefined,
15471568
scheduleDate: undefined,
15481569
},
1549-
agent: 'auto', // Add the missing 'agent' property
1570+
agent: ['auto'], // Add the missing 'agent' property
15501571
} as SessionLauncherFormData,
15511572
formValue,
15521573
);

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -821,6 +821,9 @@ class Client {
821821
if (this.isManagerVersionCompatibleWith('25.12.0')) {
822822
this._features['reservoir'] = true;
823823
}
824+
if (this.isManagerVersionCompatibleWith('25.15.0')) {
825+
this._features['multi-agents'] = true;
826+
}
824827
}
825828

826829
/**

0 commit comments

Comments
 (0)