Skip to content

Commit 3322d20

Browse files
yomybabyyomybaby
committed
feat(FR-1562): Setup Relay Subscription Environment
1 parent 72a98ab commit 3322d20

File tree

8 files changed

+1597
-551
lines changed

8 files changed

+1597
-551
lines changed

‎data/schema.graphql‎

Lines changed: 1396 additions & 245 deletions
Large diffs are not rendered by default.

‎pnpm-lock.yaml‎

Lines changed: 13 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎react/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
"fast-deep-equal": "^3.1.3",
3838
"gpt-tokenizer": "^3.2.0",
3939
"graphql": "^16.12.0",
40+
"graphql-sse": "^2.6.0",
4041
"i18next": "^25.6.0",
4142
"i18next-http-backend": "^3.0.2",
4243
"jotai": "^2.15.0",

‎react/src/RelayEnvironment.ts‎

Lines changed: 66 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,22 @@
11
// import { createClient } from "graphql-ws";
22
import { manipulateGraphQLQueryWithClientDirectives } from './helper/graphql-transformer';
3+
// import { createClient } from 'graphql-ws';
4+
import { createClient } from 'graphql-sse';
35
import {
46
Environment,
57
Network,
68
RecordSource,
79
Store,
810
FetchFunction,
9-
SubscribeFunction,
1011
RelayFeatureFlags,
12+
RequestParameters,
13+
Variables,
14+
Observable,
1115
} from 'relay-runtime';
1216

1317
RelayFeatureFlags.ENABLE_RELAY_RESOLVERS = true;
1418

15-
const fetchFn: FetchFunction = async (
16-
request,
17-
variables,
18-
// cacheConfig,
19-
// uploadables
20-
) => {
19+
const waitForBAIClient = async () => {
2120
//@ts-ignore
2221
if (globalThis.backendaiclient === undefined) {
2322
// If globalThis.backendaiclient is not defined, wait for the backend-ai-connected event.
@@ -33,6 +32,24 @@ const fetchFn: FetchFunction = async (
3332
document.addEventListener('backend-ai-connected', onBackendAIConnected);
3433
});
3534
}
35+
};
36+
const getSubscriptionEndpoint = async () => {
37+
await waitForBAIClient();
38+
let api_endpoint: any = localStorage.getItem('backendaiwebui.api_endpoint');
39+
if (api_endpoint !== null) {
40+
api_endpoint = api_endpoint.replace(/^"+|"+$/g, ''); // Remove trailing slashes
41+
api_endpoint += '/func/admin/gql';
42+
}
43+
return api_endpoint;
44+
};
45+
46+
const fetchFn: FetchFunction = async (
47+
request,
48+
variables,
49+
// cacheConfig,
50+
// uploadables
51+
) => {
52+
await waitForBAIClient();
3653

3754
const transformedQuery = manipulateGraphQLQueryWithClientDirectives(
3855
request.text || '',
@@ -72,37 +89,53 @@ const fetchFn: FetchFunction = async (
7289
return result;
7390
};
7491

75-
const subscribeFn: SubscribeFunction | undefined = undefined;
76-
77-
// if (typeof window !== "undefined") {
78-
// // We only want to setup subscriptions if we are on the client.
79-
// const subscriptionsClient = createClient({
80-
// url: WEBSOCKET_ENDPOINT,
81-
// });
92+
const subscriptionsClient = createClient({
93+
url: getSubscriptionEndpoint,
94+
headers: () => {
95+
const sessionId: string | undefined =
96+
// @ts-ignore
97+
globalThis.backendaiclient?._loginSessionId;
98+
const headers: Record<string, string> = {};
99+
if (sessionId) {
100+
headers['X-BackendAI-SessionID'] = sessionId;
101+
}
102+
return headers;
103+
},
104+
});
82105

83-
// subscribeFn = (request, variables) => {
84-
// // To understand why we return Observable.create<any>,
85-
// // please see: https://github.com/enisdenjo/graphql-ws/issues/316#issuecomment-1047605774
86-
// return Observable.create<any>((sink) => {
87-
// if (!request.text) {
88-
// return sink.error(new Error("Operation text cannot be empty"));
89-
// }
106+
function fetchForSubscribe(
107+
operation: RequestParameters,
108+
variables: Variables,
109+
): Observable<any> {
110+
return Observable.create((sink) => {
111+
if (!operation.text) {
112+
return sink.error(new Error('Operation text cannot be empty'));
113+
}
114+
const transformedOperation = manipulateGraphQLQueryWithClientDirectives(
115+
operation.text || '',
116+
variables,
117+
(version) => {
118+
// @ts-ignore
119+
return !globalThis.backendaiclient?.isManagerVersionCompatibleWith(
120+
version,
121+
);
122+
},
123+
);
90124

91-
// return subscriptionsClient.subscribe(
92-
// {
93-
// operationName: request.name,
94-
// query: request.text,
95-
// variables,
96-
// },
97-
// sink
98-
// );
99-
// });
100-
// };
101-
// }
125+
return subscriptionsClient.subscribe(
126+
{
127+
operationName: operation.name,
128+
query: transformedOperation,
129+
variables,
130+
},
131+
sink,
132+
);
133+
});
134+
}
102135

103136
function createRelayEnvironment() {
104137
return new Environment({
105-
network: Network.create(fetchFn, subscribeFn),
138+
network: Network.create(fetchFn, fetchForSubscribe),
106139
store: new Store(new RecordSource()),
107140
});
108141
}

‎react/src/components/BAIComputeSessionNodeNotificationItem.tsx‎

Lines changed: 98 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,12 @@ import {
1212
graphql,
1313
useFragment,
1414
useRelayEnvironment,
15+
useSubscription,
1516
} from 'react-relay';
1617
import { useNavigate } from 'react-router-dom';
1718
import { BAIComputeSessionNodeNotificationItemFragment$key } from 'src/__generated__/BAIComputeSessionNodeNotificationItemFragment.graphql';
1819
import { BAIComputeSessionNodeNotificationItemRefreshQuery } from 'src/__generated__/BAIComputeSessionNodeNotificationItemRefreshQuery.graphql';
20+
import { useSuspendedBackendaiClient } from 'src/hooks';
1921
import {
2022
NotificationState,
2123
useSetBAINotification,
@@ -35,6 +37,7 @@ const BAIComputeSessionNodeNotificationItem: React.FC<
3537
const { closeNotification } = useSetBAINotification();
3638
const { t } = useTranslation();
3739
const navigate = useNavigate();
40+
const baiClient = useSuspendedBackendaiClient();
3841
const node = useFragment(
3942
graphql`
4043
fragment BAIComputeSessionNodeNotificationItemFragment on ComputeSessionNode {
@@ -49,24 +52,6 @@ const BAIComputeSessionNodeNotificationItem: React.FC<
4952
sessionFrgmt,
5053
);
5154

52-
// TODO: delete this when Status subscription is implemented
53-
const [delay, setDelay] = useState<number | null>(null);
54-
UNSAFE_useAutoRefreshInterval(node?.id || '', delay);
55-
useEffect(() => {
56-
if (
57-
!node?.status ||
58-
node?.status === 'TERMINATED' ||
59-
node?.status === 'CANCELLED'
60-
) {
61-
setDelay(null);
62-
} else if (node?.status === 'RUNNING') {
63-
setDelay(15000);
64-
} else {
65-
setDelay(3000);
66-
}
67-
}, [node?.status]);
68-
// ---
69-
7055
useUpdateEffect(() => {
7156
if (node?.status === 'TERMINATED' || node?.status === 'CANCELLED') {
7257
setTimeout(() => {
@@ -77,46 +62,56 @@ const BAIComputeSessionNodeNotificationItem: React.FC<
7762

7863
return (
7964
node && (
80-
<BAINotificationItem
81-
title={
82-
<BAIText ellipsis>
83-
{t('general.Session')}:&nbsp;
84-
<BAILink
85-
style={{
86-
fontWeight: 'normal',
87-
}}
88-
title={node.name || ''}
89-
onClick={() => {
90-
navigate(
91-
`/session${node.row_id ? `?${new URLSearchParams({ sessionDetail: node.row_id }).toString()}` : ''}`,
92-
);
93-
closeNotification(notification.key);
94-
}}
95-
>
96-
{node.name}
97-
</BAILink>
98-
</BAIText>
99-
}
100-
description={
101-
<BAIFlex justify="between">
102-
<SessionStatusTag
103-
sessionFrgmt={node || null}
104-
showQueuePosition={false}
105-
showTooltip={false}
106-
/>
107-
<SessionActionButtons
108-
compact
109-
size="small"
110-
sessionFrgmt={node || null}
111-
hiddenButtonKeys={['containerCommit']}
112-
primaryAppOption={primaryAppOption}
113-
/>
114-
</BAIFlex>
115-
}
116-
footer={
117-
showDate ? dayjs(notification.created).format('lll') : undefined
118-
}
119-
/>
65+
<>
66+
<BAINotificationItem
67+
title={
68+
<BAIText ellipsis>
69+
{t('general.Session')}:&nbsp;
70+
<BAILink
71+
style={{
72+
fontWeight: 'normal',
73+
}}
74+
title={node.name || ''}
75+
onClick={() => {
76+
navigate(
77+
`/session${node.row_id ? `?${new URLSearchParams({ sessionDetail: node.row_id }).toString()}` : ''}`,
78+
);
79+
closeNotification(notification.key);
80+
}}
81+
>
82+
{node.name}
83+
</BAILink>
84+
</BAIText>
85+
}
86+
description={
87+
<BAIFlex justify="between">
88+
<SessionStatusTag
89+
sessionFrgmt={node || null}
90+
showQueuePosition={false}
91+
showTooltip={false}
92+
/>
93+
<SessionActionButtons
94+
compact
95+
size="small"
96+
sessionFrgmt={node || null}
97+
hiddenButtonKeys={['containerCommit']}
98+
primaryAppOption={primaryAppOption}
99+
/>
100+
</BAIFlex>
101+
}
102+
footer={
103+
showDate ? dayjs(notification.created).format('lll') : undefined
104+
}
105+
/>
106+
{baiClient.isManagerVersionCompatibleWith('25.16.0') && node.row_id ? (
107+
<SessionStatusRefresherUsingSubscription sessionRowId={node.row_id} />
108+
) : node.row_id && node.status ? (
109+
<UNSAFE_SessionStatusRefresher
110+
id={node.row_id}
111+
status={node.status}
112+
/>
113+
) : null}
114+
</>
120115
)
121116
);
122117
};
@@ -146,3 +141,48 @@ const UNSAFE_useAutoRefreshInterval = (
146141
).toPromise();
147142
}, delay);
148143
};
144+
145+
const SessionStatusRefresherUsingSubscription: React.FC<{
146+
sessionRowId?: string;
147+
}> = ({ sessionRowId }) => {
148+
useSubscription({
149+
subscription: graphql`
150+
subscription BAIComputeSessionNodeNotificationItemSubscription(
151+
$session_id: ID!
152+
) {
153+
schedulingEventsBySession(sessionId: $session_id) {
154+
reason
155+
session {
156+
status
157+
...SessionNodesFragment
158+
...SessionDetailContentFragment
159+
...SessionActionButtonsFragment
160+
}
161+
}
162+
}
163+
`,
164+
variables: { session_id: sessionRowId },
165+
});
166+
return null;
167+
};
168+
169+
const UNSAFE_SessionStatusRefresher: React.FC<{
170+
id: string;
171+
status: string;
172+
}> = ({ id, status }) => {
173+
// TODO: delete this when Status subscription is implemented
174+
const [delay, setDelay] = useState<number | null>(null);
175+
UNSAFE_useAutoRefreshInterval(id || '', delay);
176+
useEffect(() => {
177+
if (!status || status === 'TERMINATED' || status === 'CANCELLED') {
178+
setDelay(null);
179+
} else if (status === 'RUNNING') {
180+
setDelay(15000);
181+
} else {
182+
setDelay(3000);
183+
}
184+
}, [status]);
185+
// ---
186+
187+
return null;
188+
};

0 commit comments

Comments
 (0)