Skip to content

Commit 589e38d

Browse files
committed
feat(FR-1448): allow user to select multi agents when creating multi-node session
1 parent 9dce65a commit 589e38d

File tree

4 files changed

+60
-28
lines changed

4 files changed

+60
-28
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: 12 additions & 15 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 &
@@ -1229,14 +1229,15 @@ const ResourceAllocationFormItems: React.FC<
12291229
<AgentSelect
12301230
resourceGroup={currentResourceGroupInForm}
12311231
fetchKey={agentFetchKey}
1232-
onChange={(value, option) => {
1233-
if (value !== 'auto') {
1234-
form.setFieldsValue({
1235-
cluster_mode: 'single-node',
1236-
cluster_size: 1,
1237-
});
1238-
}
1239-
// TODO: set cluster mode to single node and cluster size to 1 when agent value is not "auto"
1232+
{...{
1233+
mode: baiClient.supports('multi-agents')
1234+
? 'multiple'
1235+
: undefined,
1236+
labelRender: baiClient.supports('multi-agents')
1237+
? ({ label, value }) => {
1238+
return value === 'auto' ? label : value;
1239+
}
1240+
: undefined,
12401241
}}
12411242
></AgentSelect>
12421243
</Form.Item>
@@ -1282,7 +1283,6 @@ const ResourceAllocationFormItems: React.FC<
12821283
onChange={(e) => {
12831284
form.validateFields().catch(() => {});
12841285
}}
1285-
disabled={getFieldValue('agent') !== 'auto'}
12861286
>
12871287
<Radio.Button value="single-node">
12881288
{t('session.launcher.SingleNode')}
@@ -1345,10 +1345,7 @@ const ResourceAllocationFormItems: React.FC<
13451345
? derivedClusterSizeMaxLimit
13461346
: undefined
13471347
}
1348-
disabled={
1349-
derivedClusterSizeMaxLimit === 1 ||
1350-
getFieldValue('agent') !== 'auto'
1351-
}
1348+
disabled={derivedClusterSizeMaxLimit === 1}
13521349
sliderProps={{
13531350
marks: {
13541351
1: '1',

react/src/pages/SessionLauncherPage.tsx

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -544,16 +544,39 @@ const SessionLauncherPage = () => {
544544
preopen_ports: transformPortValuesToNumbers(values.ports),
545545

546546
// Agent selection (optional)
547-
...(baiClient.supports('agent-select') &&
548-
!baiClient?._config?.hideAgents &&
549-
values.agent !== 'auto'
550-
? {
551-
// Filter out undefined values
552-
agent_list: [values.agent].filter(
553-
(agent): agent is string => !!agent,
547+
...(() => {
548+
if (values.agent === undefined) return {};
549+
550+
const agents = Array.isArray(values.agent)
551+
? values.agent
552+
: [values.agent];
553+
554+
if (
555+
!baiClient.supports('agent-select') ||
556+
baiClient?._config?.hideAgents ||
557+
_.isEqual(agents, ['auto']) ||
558+
agents.length === 0
559+
) {
560+
return {};
561+
}
562+
563+
if (
564+
!baiClient.supports('multi-agents') &&
565+
values.cluster_mode === 'multi-node'
566+
) {
567+
// The server now requires agents equivalent to the cluster size.
568+
return {
569+
agent_list: Array.from(
570+
{ length: values.cluster_size },
571+
() => agents[0],
554572
),
555-
}
556-
: undefined),
573+
};
574+
}
575+
576+
return {
577+
agent_list: _.filter(agents, (item) => !!item),
578+
};
579+
})(),
557580
},
558581
},
559582
};
@@ -1606,7 +1629,7 @@ const SessionLauncherPage = () => {
16061629
command: undefined,
16071630
scheduleDate: undefined,
16081631
},
1609-
agent: 'auto', // Add the missing 'agent' property
1632+
agent: ['auto'], // Add the missing 'agent' property
16101633
} as SessionLauncherFormData,
16111634
formValue,
16121635
);

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -813,6 +813,9 @@ class Client {
813813
if (this.isManagerVersionCompatibleWith('25.13.0')) {
814814
this._features['pending-session-list'] = true;
815815
}
816+
if (this.isManagerVersionCompatibleWith('25.15.0')) {
817+
this._features['multi-agents'] = true;
818+
}
816819
}
817820

818821
/**

0 commit comments

Comments
 (0)