Skip to content

Commit f70afb7

Browse files
larslllarsll
authored
Rework model load mechanism to avoid selecting model when none is loaded (#76)
1 parent ba2a7bc commit f70afb7

File tree

2 files changed

+134
-132
lines changed

2 files changed

+134
-132
lines changed

‎website/src/common/hooks/use-models.ts‎

Lines changed: 102 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ interface ModelsContextState {
4343
// Add flashbar items to the context
4444
modelFlashbarItems: FlashbarProps.MessageDefinition[];
4545
clearModelFlashbar: () => void;
46+
checkModelLoadStatus: () => Promise<boolean>;
4647
}
4748

4849
export const ModelsContext = createContext<ModelsContextState | null>(null);
@@ -82,43 +83,46 @@ export const useModelsProvider = () => {
8283
}, []);
8384

8485
// Function to fetch available models
85-
const fetchModels = useCallback(async () => {
86-
// Don't fetch if not authenticated
87-
if (!isAuthenticated) {
88-
return [];
89-
}
86+
const fetchModels = useCallback(
87+
async (isModelLoaded: boolean) => {
88+
// Don't fetch if not authenticated
89+
if (!isAuthenticated) {
90+
return [];
91+
}
9092

91-
try {
92-
const response = await ApiHelper.get<ModelsResponse>("models");
93-
if (response) {
94-
const options = response.models.map((model) => ({
95-
label: model.model_folder_name,
96-
value: model.model_folder_name,
97-
description: model.model_sensors.join(", "),
98-
disabled: model.is_select_disabled,
99-
}));
100-
setModelOptions(options);
101-
102-
// Restore previously selected model if it exists in the new list
103-
const selectedModelName = localStorage.getItem("selectedModelName");
104-
if (selectedModelName) {
105-
const previousModel = options.find((model) => model.value === selectedModelName);
106-
if (previousModel && isModelLoaded) {
107-
setSelectedModel(previousModel);
108-
} else {
109-
// If the previous model is not found, clear the selection
110-
setSelectedModel(null);
93+
try {
94+
const response = await ApiHelper.get<ModelsResponse>("models");
95+
if (response) {
96+
const options = response.models.map((model) => ({
97+
label: model.model_folder_name,
98+
value: model.model_folder_name,
99+
description: model.model_sensors.join(", "),
100+
disabled: model.is_select_disabled,
101+
}));
102+
setModelOptions(options);
103+
104+
// Restore previously selected model if it exists in the new list
105+
const selectedModelName = sessionStorage.getItem("selectedModelName");
106+
if (selectedModelName) {
107+
const previousModel = options.find((model) => model.value === selectedModelName);
108+
if (previousModel && isModelLoaded) {
109+
setSelectedModel(previousModel);
110+
} else {
111+
// If the previous model is not found, clear the selection
112+
setSelectedModel(null);
113+
}
111114
}
112-
}
113115

114-
return options;
116+
return options;
117+
}
118+
return [];
119+
} catch (error) {
120+
console.error("Error fetching models:", error);
121+
return [];
115122
}
116-
return [];
117-
} catch (error) {
118-
console.error("Error fetching models:", error);
119-
return [];
120-
}
121-
}, [isAuthenticated, isModelLoaded]);
123+
},
124+
[isAuthenticated]
125+
);
122126

123127
// Function to check if a model is currently loaded
124128
const checkModelLoadStatus = useCallback(async () => {
@@ -140,103 +144,92 @@ export const useModelsProvider = () => {
140144
} else {
141145
setIsModelLoaded(false);
142146
setIsModelLoading(false);
147+
setSelectedModel(null); // Clear selected model if loading failed
143148
return false;
144149
}
145150
} catch (error) {
146151
console.error("Error checking model status:", error);
147152
setIsModelLoaded(false);
148153
setIsModelLoading(false);
154+
setSelectedModel(null); // Clear selected model if loading failed
149155
return false;
150156
}
151157
}, [isAuthenticated]);
152158

153159
// Function to poll for model loading status with retry limit
154-
const pollModelLoadingStatus = useCallback(
155-
(retryCount = 0) => {
156-
// Don't poll if not authenticated
157-
if (!isAuthenticated) {
158-
return;
159-
}
160-
161-
// Maximum number of retries (30 seconds with 1-second interval)
162-
const MAX_RETRIES = 30;
160+
const pollModelLoadingStatus = useCallback((retryCount = 0) => {
161+
// Maximum number of retries (30 seconds with 1-second interval)
162+
const MAX_RETRIES = 30;
163163

164-
if (retryCount >= MAX_RETRIES) {
165-
// Too many retries, stop polling
166-
setLoadStatus({
167-
loading: false,
168-
success: false,
169-
error: "Model loading timed out",
170-
});
171-
return;
172-
}
164+
if (retryCount >= MAX_RETRIES) {
165+
// Too many retries, stop polling
166+
setLoadStatus({
167+
loading: false,
168+
success: false,
169+
error: "Model loading timed out",
170+
});
171+
return;
172+
}
173173

174-
// Check model status
175-
ApiHelper.get<ModelLoadingResponse>("isModelLoading")
176-
.then((response) => {
177-
if (response?.success) {
178-
if (response.isModelLoading === "loaded") {
179-
// Model loaded successfully
180-
setIsModelLoaded(true);
181-
setIsModelLoading(false);
182-
setLoadStatus({
183-
loading: false,
184-
success: true,
185-
error: null,
186-
});
187-
} else if (response.isModelLoading === "loading") {
188-
// Still loading, continue polling
189-
setIsModelLoaded(false);
190-
setIsModelLoading(true);
191-
setTimeout(() => pollModelLoadingStatus(retryCount + 1), 1000);
192-
} else if (response.isModelLoading === "failed") {
193-
// Loading failed
194-
setIsModelLoaded(false);
195-
setIsModelLoading(false);
196-
setLoadStatus({
197-
loading: false,
198-
success: false,
199-
error: "Model loading failed",
200-
});
201-
} else {
202-
// Unknown state, try again
203-
setIsModelLoaded(false);
204-
setIsModelLoading(false);
205-
setTimeout(() => pollModelLoadingStatus(retryCount + 1), 1000);
206-
}
207-
} else {
208-
// API call failed
174+
// Check model status
175+
ApiHelper.get<ModelLoadingResponse>("isModelLoading")
176+
.then((response) => {
177+
if (response?.success) {
178+
if (response.isModelLoading === "loaded") {
179+
// Model loaded successfully
180+
setIsModelLoaded(true);
181+
setIsModelLoading(false);
182+
setLoadStatus({
183+
loading: false,
184+
success: true,
185+
error: null,
186+
});
187+
} else if (response.isModelLoading === "loading") {
188+
// Still loading, continue polling
189+
setIsModelLoaded(false);
190+
setIsModelLoading(true);
191+
setTimeout(() => pollModelLoadingStatus(retryCount + 1), 1000);
192+
} else if (response.isModelLoading === "failed" || response.isModelLoading === "error") {
193+
// Loading failed
209194
setIsModelLoaded(false);
210195
setIsModelLoading(false);
211196
setLoadStatus({
212197
loading: false,
213198
success: false,
214-
error: "Error checking model loading status",
199+
error: "Model loading failed",
215200
});
201+
} else {
202+
// Unknown state, try again
203+
setIsModelLoaded(false);
204+
setIsModelLoading(false);
205+
setTimeout(() => pollModelLoadingStatus(retryCount + 1), 1000);
216206
}
217-
})
218-
.catch((error) => {
219-
console.error("Error polling model status:", error);
207+
} else {
208+
// API call failed
220209
setIsModelLoaded(false);
221210
setIsModelLoading(false);
222211
setLoadStatus({
223212
loading: false,
224213
success: false,
225-
error: error instanceof Error ? error.message : "Error checking model status",
214+
error: "Error checking model loading status",
226215
});
216+
}
217+
})
218+
.catch((error) => {
219+
console.error("Error polling model status:", error);
220+
setIsModelLoaded(false);
221+
setIsModelLoading(false);
222+
setLoadStatus({
223+
loading: false,
224+
success: false,
225+
error: error instanceof Error ? error.message : "Error checking model status",
227226
});
228-
},
229-
[isAuthenticated]
230-
);
227+
});
228+
}, []);
231229

232230
// Function to load a model
233231
const loadModel = useCallback(
234232
async (modelName: string) => {
235-
// Don't load if not authenticated
236-
if (!isAuthenticated) {
237-
return false;
238-
}
239-
240233
try {
241234
// Reset states
242235
setLoadStatus({
@@ -253,7 +246,7 @@ export const useModelsProvider = () => {
253246
);
254247

255248
if (modelResponse?.success) {
256-
localStorage.setItem("selectedModelName", modelName);
249+
sessionStorage.setItem("selectedModelName", modelName);
257250

258251
// Start polling with initial retry count of 0
259252
pollModelLoadingStatus(0);
@@ -278,19 +271,14 @@ export const useModelsProvider = () => {
278271
return false;
279272
}
280273
},
281-
[pollModelLoadingStatus, isAuthenticated]
274+
[pollModelLoadingStatus]
282275
);
283276

284277
// Function to reload models list
285278
const reloadModels = useCallback(async () => {
286-
// Don't reload if not authenticated
287-
if (!isAuthenticated) {
288-
return;
289-
}
290-
291-
await fetchModels();
292-
await checkModelLoadStatus();
293-
}, [fetchModels, checkModelLoadStatus, isAuthenticated]);
279+
const isModelLoaded = await checkModelLoadStatus();
280+
await fetchModels(isModelLoaded);
281+
}, [fetchModels, checkModelLoadStatus]);
294282

295283
// Initialize models and check status on mount
296284
useEffect(() => {
@@ -311,11 +299,12 @@ export const useModelsProvider = () => {
311299
}
312300

313301
const initialize = async () => {
314-
const options = await fetchModels();
315-
await checkModelLoadStatus();
302+
console.debug("Initializing models...");
303+
const isModelLoaded = await checkModelLoadStatus();
304+
const options = await fetchModels(isModelLoaded);
316305

317306
// Restore previously selected model
318-
const selectedModelName = localStorage.getItem("selectedModelName");
307+
const selectedModelName = sessionStorage.getItem("selectedModelName");
319308
if (selectedModelName && options.length > 0) {
320309
const previousModel = options.find((model) => model.value === selectedModelName);
321310
if (previousModel) {
@@ -381,6 +370,7 @@ export const useModelsProvider = () => {
381370
loadStatus,
382371
modelFlashbarItems,
383372
clearModelFlashbar,
373+
checkModelLoadStatus,
384374
};
385375

386376
return contextValue;

0 commit comments

Comments
 (0)