gpustack-ui/src/pages/llmodels/deployments.tsx

import { PaginationKey } from '@/config/settings';
import useSetChunkRequest from '@/hooks/use-chunk-request';
import { usePaginationStatus } from '@/hooks/use-pagination-status';
import { useTableMultiSort } from '@/hooks/use-table-sort';
import useUpdateChunkedList from '@/hooks/use-update-chunk-list';
import { TableOrder, TableProvider } from '@gpustack/core-ui';
import { useMemoizedFn } from 'ahooks';
import _ from 'lodash';
import qs from 'query-string';
import {
  forwardRef,
  useEffect,
  useImperativeHandle,
  useRef,
  useState
} from 'react';
import {
  MODELS_API,
  MODEL_INSTANCE_API,
  queryModelsInstances,
  queryModelsList
} from './apis';
import TableList from './components/table-list';
import { ListItem } from './config/types';

const Models = forwardRef((props, ref) => {
  const { pagination, setPagination } = usePaginationStatus(
    PaginationKey.Deployments
  );
  const { sortOrder, handleMultiSortChange } = useTableMultiSort();
  const { setChunkRequest, createAxiosToken } = useSetChunkRequest();
  const { setChunkRequest: setModelInstanceChunkRequest } =
    useSetChunkRequest();
  const [modelInstances, setModelInstances] = useState<any[]>([]);
  const [dataSource, setDataSource] = useState<{
    dataList: ListItem[];
    deletedIds: number[];
    loading: boolean;
    loadend: boolean;
    total: number;
  }>({
    dataList: [],
    deletedIds: [],
    loading: false,
    loadend: false,
    total: 0
  });
  const chunkRequedtRef = useRef<any>();
  const chunkInstanceRequedtRef = useRef<any>();
  const isPageHidden = useRef(false);
  const instancesToken = useRef<any>();
  let axiosToken = createAxiosToken();
  const [queryParams, setQueryParams] = useState({
    page: 1,
    perPage: 10,
    search: '',
    cluster_id: 0,
    categories: [],
    state: '',
    sort_by: '',
    ...pagination
  });

  const { updateChunkedList, cacheDataListRef } = useUpdateChunkedList({
    events: ['UPDATE'],
    dataList: dataSource.dataList,
    setDataList(list, opts?: any) {
      setDataSource((pre) => {
        return {
          total: pre.total,
          loading: false,
          loadend: true,
          dataList: list,
          deletedIds: opts?.deletedIds || []
        };
      });
    }
  });

  const {
    updateChunkedList: updateInstanceChunkedList,
    cacheDataListRef: cacheInsDataListRef
  } = useUpdateChunkedList({
    dataList: modelInstances,
    limit: 100,
    setDataList: setModelInstances
  });

  const [filterValues, setFilterValues] = useState<any>({});

  const getAllModelInstances = useMemoizedFn(async () => {
    try {
      instancesToken.current?.cancel?.();
      instancesToken.current = createAxiosToken();
      const params = {
        page: -1
      };
      const res: any = await queryModelsInstances(params, {
        token: instancesToken.current.token
      });
      cacheInsDataListRef.current = res.items || [];
      setModelInstances(res.items || []);
    } catch (error) {
      // ignore
    }
  });

  const fetchData = useMemoizedFn(
    async (params?: {
      loadingVal?: boolean;
      query?: {
        page: number;
        perPage: number;
        search: string;
        categories: any[];
        sort_by: string;
      };
    }) => {
      const { loadingVal, query } = params || {};
      axiosToken?.cancel?.();
      axiosToken = createAxiosToken();
      setDataSource((pre) => {
        pre.loading = loadingVal ?? true;
        return { ...pre };
      });
      try {
        const params = {
          ..._.pickBy(query || queryParams, (val: any) => !!val)
        };
        const res: any = await queryModelsList(params, {
          cancelToken: axiosToken.token
        });

        // if the current page is beyond total page, fetch again
        if (
          !res.items.length &&
          params.page > res.pagination.totalPage &&
          res.pagination.totalPage > 0
        ) {
          const newParams = {
            ...params,
            page: res.pagination.totalPage
          };
          const newRes: any = await queryModelsList(newParams, {
            cancelToken: axiosToken.token
          });
          setDataSource({
            dataList: newRes.items || [],
            loading: false,
            loadend: true,
            total: newRes.pagination.total,
            deletedIds: []
          });
          return;
        }

        setDataSource({
          dataList: res.items || [],
          loading: false,
          loadend: true,
          total: res.pagination.total,
          deletedIds: []
        });
      } catch (error) {
        if (!isPageHidden.current) {
          setDataSource({
            dataList: [],
            loading: false,
            loadend: true,
            total: dataSource.total,
            deletedIds: []
          });
        }
      }
    }
  );

  const handleQueryChange = (params: any) => {
    setQueryParams({
      ...queryParams,
      ...params
    });
    fetchData({ query: { ...queryParams, ...params } });
  };

  const handlePageChange = useMemoizedFn(
    (page: number, pageSize: number | undefined) => {
      console.log('page change', page, pageSize);
      handleQueryChange({
        page: page,
        perPage: pageSize || 10
      });
      setPagination({
        perPage: pageSize || 10
      });
    }
  );

  const updateHandler = (list: any) => {
    _.each(list, (data: any) => {
      updateChunkedList(data);
    });
  };

  const updateInstanceHandler = (list: any) => {
    // filter the data
    _.each(list, (data: any) => {
      updateInstanceChunkedList(data);
    });
  };

  const createModelsChunkRequest = useMemoizedFn(
    async (params?: {
      search: string;
      categories: any[];
      cluster_id: number;
      state?: string;
    }) => {
      const search = params?.search || queryParams.search;
      const categories = params?.categories || queryParams.categories;
      chunkRequedtRef.current?.current?.cancel?.();
      try {
        const query = {
          search: search,
          categories: categories
        };
        chunkRequedtRef.current = setChunkRequest({
          url: `${MODELS_API}?${qs.stringify(_.pickBy(query, (val: any) => !!val))}`,
          handler: updateHandler
        });
      } catch (error) {
        // ignore
      }
    }
  );

  const createModelsInstanceChunkRequest = useMemoizedFn(async () => {
    chunkInstanceRequedtRef.current?.current?.cancel?.();
    cacheInsDataListRef.current = [];
    try {
      chunkInstanceRequedtRef.current = setModelInstanceChunkRequest({
        url: `${MODEL_INSTANCE_API}`,
        params: {},
        handler: updateInstanceHandler
      });
    } catch (error) {
      // ignore
    }
  });

  const cancelRequestsOnPageInactive = useMemoizedFn(() => {
    isPageHidden.current = true;
    chunkRequedtRef.current?.current?.cancel?.();
    cacheDataListRef.current = [];
    cacheInsDataListRef.current = [];
    chunkInstanceRequedtRef.current?.current?.cancel?.();
    instancesToken.current?.cancel?.();
  });

  const resumeRequestsOnPageActive = useMemoizedFn(async () => {
    isPageHidden.current = false;
    fetchData({
      loadingVal: false
    });
    await getAllModelInstances();
    await createModelsInstanceChunkRequest();
    await createModelsChunkRequest();
  });

  const handleOnCancelViewLogs = useMemoizedFn(async () => {
    resumeRequestsOnPageActive();
  });

  const handleSearchBySilent = useMemoizedFn(async () => {
    await new Promise((resolve) => {
      setTimeout(resolve, 300);
    });
    fetchData({
      loadingVal: false
    });
    // re-align the instance cache with the backend as a backstop, in case a
    // DELETE watch event for terminated instances was missed
    getAllModelInstances();
  });

  // proactively drop a model's instances from the cache when it is stopped, so
  // a stale (terminating) instance can't linger and show up alongside the new
  // one after an immediate restart
  const handleStop = useMemoizedFn(async (modelIds: number[]) => {
    const idSet = new Set(modelIds);
    cacheInsDataListRef.current = cacheInsDataListRef.current.filter(
      (item) => !idSet.has(item.model_id)
    );
    setModelInstances(cacheInsDataListRef.current);
    handleSearchBySilent();
  });

  const handleSearch = useMemoizedFn(async (params?: any) => {
    await fetchData(params);
  });

  const debounceUpdateFilter = _.debounce((e: any) => {
    handleQueryChange({
      page: 1,
      search: e.target.value
    });
    createModelsChunkRequest({
      search: e.target.value,
      categories: queryParams.categories,
      cluster_id: queryParams.cluster_id || 0
    });
  }, 350);

  const handleNameChange = useMemoizedFn(debounceUpdateFilter);

  const handleCategoryChange = async (value: any) => {
    handleQueryChange({
      page: 1,
      categories: value
    });
    createModelsChunkRequest({
      search: queryParams.search,
      cluster_id: queryParams.cluster_id || 0,
      categories: value
    });
  };

  const handleOnStatusChange = async (value: string | undefined) => {
    handleQueryChange({
      page: 1,
      state: value
    });
    createModelsChunkRequest({
      search: queryParams.search,
      categories: queryParams.categories,
      cluster_id: queryParams.cluster_id || 0,
      state: value
    });
  };

  const handleClusterChange = async (value: number) => {
    handleQueryChange({
      page: 1,
      cluster_id: value
    });
    createModelsChunkRequest({
      search: queryParams.search,
      categories: queryParams.categories,
      cluster_id: value
    });
  };

  const handleOnSortChange = (order: TableOrder | Array<TableOrder>) => {
    const orderList = Array.isArray(order) ? order : [order];
    if (orderList[0].columnKey === 'replicas') {
      orderList.push({
        columnKey: 'ready_replicas',
        order: orderList[0].order
      });
    }
    const sortKeys = handleMultiSortChange(orderList);
    setQueryParams((pre: any) => {
      return {
        ...pre,
        page: 1,
        sort_by: sortKeys.join(',')
      };
    });
    fetchData({
      query: {
        ...queryParams,
        page: 1,
        sort_by: sortKeys.join(',')
      }
    });
  };

  const handleOnFilterChange = (filters: any) => {
    handleQueryChange({
      page: 1,
      ...filters
    });
    createModelsChunkRequest({
      search: queryParams.search,
      ...filters
    });

    setFilterValues(filters);
  };

  const handleDeleteInstanceFromCache = (id: number) => {
    cacheInsDataListRef.current = cacheInsDataListRef.current.filter(
      (item) => item.id !== id
    );
    setModelInstances(cacheInsDataListRef.current);
  };

  useEffect(() => {
    let timer: any = null;
    // fetch data first time
    const getTableData = async (loadingVal?: boolean) => {
      axiosToken?.cancel?.();
      axiosToken = createAxiosToken();
      setDataSource((pre) => {
        pre.loading = loadingVal ?? true;
        return { ...pre };
      });
      try {
        const params = {
          ..._.pickBy(queryParams, (val: any) => !!val)
        };
        const res: any = await queryModelsList(params, {
          cancelToken: axiosToken.token
        });
        return res;
      } catch (error) {
        return {};
      }
    };

    const init = async () => {
      const [modelRes] = await Promise.all([getTableData()]);

      setDataSource({
        dataList: modelRes.items || [],
        loading: false,
        loadend: true,
        total: modelRes.pagination?.total || 0,
        deletedIds: []
      });

      clearTimeout(timer);
      timer = setTimeout(() => {
        createModelsInstanceChunkRequest();
        createModelsChunkRequest();
      }, 1000);
    };
    init();
    return () => {
      clearTimeout(timer);
      axiosToken?.cancel?.();
      cancelRequestsOnPageInactive();
    };
  }, []);

  const setDisableExpand = useMemoizedFn((record: any) => {
    return !record?.replicas;
  });

  useImperativeHandle(ref, () => ({
    cancelRequestsOnPageInactive,
    resumeRequestsOnPageActive
  }));

  return (
    <TableProvider
      value={{
        allChildren: modelInstances,
        setDisableExpand: setDisableExpand
      }}
    >
      <TableList
        dataSource={dataSource.dataList}
        onStatusChange={handleOnStatusChange}
        handleNameChange={handleNameChange}
        handleCategoryChange={handleCategoryChange}
        handleClusterChange={handleClusterChange}
        handleSearch={handleSearch}
        handlePageChange={handlePageChange}
        handleDeleteSuccess={fetchData}
        handleOnToggleExpandAll={createModelsInstanceChunkRequest}
        onViewLogs={cancelRequestsOnPageInactive}
        onCancelViewLogs={handleOnCancelViewLogs}
        onStop={handleStop}
        onStart={handleSearchBySilent}
        onTableSort={handleOnSortChange}
        onFilterChange={handleOnFilterChange}
        onDeleteInstanceFromCache={handleDeleteInstanceFromCache}
        sortOrder={sortOrder}
        queryParams={queryParams}
        loading={dataSource.loading}
        loadend={dataSource.loadend}
        total={dataSource.total}
        deleteIds={dataSource.deletedIds}
        filterValues={filterValues}
      ></TableList>
    </TableProvider>
  );
});

export default Models;