Skip to content

Commit

Permalink
feat: terminiate sessions in NEO session list
Browse files Browse the repository at this point in the history
  • Loading branch information
@yomybaby
yomybaby committed Feb 14, 2025
1 parent ffa9a19 commit a499b90
Show file tree
Hide file tree
Showing 4 changed files with 272 additions and 138 deletions.
177 changes: 92 additions & 85 deletions react/src/components/ComputeSessionNodeItems/TerminateSessionModal.tsx
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import { filterEmptyItem } from '../../helper';
import { BackendAIClient, useSuspendedBackendaiClient } from '../../hooks';
import { useCurrentUserRole } from '../../hooks/backendai';
import { useTanMutation } from '../../hooks/reactQueryAlias';
import { useSetBAINotification } from '../../hooks/useBAINotification';
import { useCurrentProjectValue } from '../../hooks/useCurrentProject';
import { usePainKiller } from '../../hooks/usePainKiller';
import { usePromiseTracker } from '../../usePromiseTracker';
import BAIModal from '../BAIModal';
import Flex from '../Flex';
import {
Expand All @@ -21,10 +22,12 @@ import { fetchQuery, useFragment, useRelayEnvironment } from 'react-relay';

interface TerminateSessionModalProps
extends Omit<ModalProps, 'onOk' | 'onCancel'> {
sessionFrgmts: TerminateSessionModalFragment$key;
sessionFrgmts?: TerminateSessionModalFragment$key;
onRequestClose: (success: boolean) => void;
}

// Cannot destroy sessions in scheduled/preparing/pulling/prepared/creating/terminating/error status

const useStyle = createStyles(({ css, token }) => {
return {
custom: css`
Expand Down Expand Up @@ -173,7 +176,7 @@ const terminateApp = async (
};

const TerminateSessionModal: React.FC<TerminateSessionModalProps> = ({
sessionFrgmts: sessionFrgmt,
sessionFrgmts,
onRequestClose,
...modalProps
}) => {
Expand All @@ -198,7 +201,7 @@ const TerminateSessionModal: React.FC<TerminateSessionModalProps> = ({
}
}
`,
sessionFrgmt,
sessionFrgmts,
);

const [isForce, setIsForce] = useState(false);
Expand All @@ -208,39 +211,40 @@ const TerminateSessionModal: React.FC<TerminateSessionModalProps> = ({

const currentProject = useCurrentProjectValue();

const terminateMutation = useTanMutation({
mutationFn: async (session: Session) => {
return terminateApp(
session,
baiClient._config.accessKey,
currentProject.id,
baiClient,
)
.catch((e) => {
return {
error: e,
};
})
.then((result) => {
const err = result?.error;
if (
err === undefined || //no error
(err && // Even if wsproxy address is invalid, session must be deleted.
err.message &&
(err.statusCode === 404 || err.statusCode === 500))
) {
// BAI client destroy try to request 3times as default
return baiClient.destroy(
session.row_id,
baiClient._config.accessKey,
isForce,
);
} else {
throw err;
}
});
},
});
const { pendingCount, trackPromise } = usePromiseTracker();

const terminiateSession = (session: Session) => {
return terminateApp(
session,
baiClient._config.accessKey,
currentProject.id,
baiClient,
)
.catch((e) => {
return {
error: e,
};
})
.then((result) => {
const err = result?.error;
if (
err === undefined || //no error
(err && // Even if wsproxy address is invalid, session must be deleted.
err.message &&
(err.statusCode === 404 || err.statusCode === 500))
) {
// BAI client destroy try to request 3times as default
return baiClient.destroy(
session.row_id,
baiClient._config.accessKey,
isForce,
);
} else {
throw err;
}
});
};

const relayEvn = useRelayEnvironment();
const painKiller = usePainKiller();
const { upsertNotification } = useSetBAINotification();
Expand All @@ -250,54 +254,57 @@ const TerminateSessionModal: React.FC<TerminateSessionModalProps> = ({
centered
title={t('session.TerminateSession')}
open={openTerminateModal}
confirmLoading={terminateMutation.isPending}
confirmLoading={pendingCount > 0}
onOk={() => {
if (sessions[0]?.row_id) {
const session = sessions[0];
terminateMutation
.mutateAsync(session)
.then(() => {
setIsForce(false);
onRequestClose(true);
})
.catch((err) => {
upsertNotification({
message: painKiller.relieve(err?.title),
description: err?.message,
open: true,
});
})
.finally(() => {
// TODO: remove below code after session list migration to React
const event = new CustomEvent(
'backend-ai-session-list-refreshed',
{
detail: 'running',
},
);
document.dispatchEvent(event);
const promises = _.map(
filterEmptyItem(_.castArray(sessions)),
(session) => {
return terminiateSession(session)
.catch((err) => {
upsertNotification({
message: painKiller.relieve(err?.title),
description: err?.message,
open: true,
});
})
.finally(() => {
// TODO: remove below code after session list migration to React
const event = new CustomEvent(
'backend-ai-session-list-refreshed',
{
detail: 'running',
},
);
document.dispatchEvent(event);

// refetch session node
return fetchQuery<TerminateSessionModalRefetchQuery>(
relayEvn,
graphql`
query TerminateSessionModalRefetchQuery(
$id: GlobalIDField!
$project_id: UUID!
) {
compute_session_node(id: $id, project_id: $project_id) {
id
status
// refetch session node
return fetchQuery<TerminateSessionModalRefetchQuery>(
relayEvn,
graphql`
query TerminateSessionModalRefetchQuery(
$id: GlobalIDField!
$project_id: UUID!
) {
compute_session_node(id: $id, project_id: $project_id) {
id
status
}
}
}
`,
{
id: session.id,
project_id: currentProject.id,
},
).toPromise();
});
}
`,
{
id: session.id,
project_id: currentProject.id,
},
).toPromise();
});
},
);
promises.map(trackPromise);
Promise.allSettled(promises).then((results) => {
setIsForce(false);

onRequestClose(true);
});
}}
okText={isForce ? t('button.ForceTerminate') : t('session.Terminate')}
okType="danger"
Expand All @@ -320,9 +327,9 @@ const TerminateSessionModal: React.FC<TerminateSessionModalProps> = ({
{t('userSettings.SessionTerminationDialog')}
</Typography.Text>
<Typography.Text mark>
{sessions.length === 1
? sessions[0]?.name
: `${sessions.length} sessions`}
{sessions?.length === 1
? sessions?.[0]?.name
: `${sessions?.length} sessions`}
</Typography.Text>
<Checkbox
checked={isForce}
Expand Down
14 changes: 10 additions & 4 deletions react/src/components/SessionNodes.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,22 @@ import SessionStatusTag from './ComputeSessionNodeItems/SessionStatusTag';
import Flex from './Flex';
import SessionDetailDrawer from './SessionDetailDrawer';
import SessionUsageMonitor from './SessionUsageMonitor';
import { SessionNodesFragment$key } from './__generated__/SessionNodesFragment.graphql';
import {
SessionNodesFragment$data,
SessionNodesFragment$key,
} from './__generated__/SessionNodesFragment.graphql';
import { TableProps, theme } from 'antd/lib';
import graphql from 'babel-plugin-relay/macro';
import React, { useState } from 'react';
import { useTranslation } from 'react-i18next';
import { useFragment } from 'react-relay';

interface SessionNodesProps extends Omit<TableProps, 'dataSource' | 'columns'> {
export type SessionNodeInList = NonNullable<SessionNodesFragment$data[number]>;
interface SessionNodesProps
extends Omit<TableProps<SessionNodeInList>, 'dataSource' | 'columns'> {
sessionsFrgmt: SessionNodesFragment$key;
}

const SessionNodes: React.FC<SessionNodesProps> = ({
sessionsFrgmt,
...tableProps
Expand All @@ -28,7 +34,7 @@ const SessionNodes: React.FC<SessionNodesProps> = ({
const sessions = useFragment(
graphql`
fragment SessionNodesFragment on ComputeSessionNode @relay(plural: true) {
id
id @required(action: NONE)
row_id @required(action: NONE)
name
...SessionStatusTagFragment
Expand All @@ -49,7 +55,7 @@ const SessionNodes: React.FC<SessionNodesProps> = ({
neoStyle
// TODO: fix type
// @ts-ignore
rowKey={(record) => record.row_id as string}
rowKey={(record) => record.id as string}
size="small"
dataSource={filteredSessions}
scroll={{ x: 'max-content' }}
Expand Down
Loading

0 comments on commit a499b90

Please sign in to comment.