Commit 4aeff5c3 authored by Hines, Jesse's avatar Hines, Jesse
Browse files

Use infinite queries for jobs and fix some bugs

parent 3f208647

src/util/hooks/useReplay.ts

+20 −7
Original line number Diff line number Diff line 
import { useEffect } from "react";

import { useQuery, UseQueryOptions, useQueryClient } from "@tanstack/react-query";

import { useQuery, UseQueryOptions, useQueryClient, useInfiniteQuery } from "@tanstack/react-query";

import { sortBy } from "lodash";

import {

  toDate,  isEqual as isDateEqual, min as minDate, addSeconds, subSeconds, differenceInSeconds,
@@ -164,6 +164,8 @@ export type UseJobReplayOptions = { 


export type UseJobReplayResult = {

  data: Job[],

  hasNextPage: boolean,

  fetchNextPage: () => void,

  maxTimestamp: Date|undefined,

  currentTimestamp: Date|undefined,

  nextTimestamp: Date|undefined,
@@ -206,7 +208,7 @@ export const useJobReplay = ({ 
    enabled = false;

  }



  const { data: response } = useQuery({

  const { data, fetchNextPage, hasNextPage, isFetching } = useInfiniteQuery({

    ...simulationSchedulerJobs(sim?.id ?? '', {

      start: queryStart?.toISOString(), end: queryEnd?.toISOString(),

      limit: 1000,
@@ -218,12 +220,14 @@ export const useJobReplay = ({ 
  });



  let jobs = (

    response?.results

    data?.pages

      .map((page) => page.results)

      .flat()

      .map(j => ({...j, state_current: computeJobState(j, currentTimestamp)}))

      // Filter out unsubmitted jobs, and completed jobs after a few steps

      .filter(j =>

        j.state_current != "UNSUBMITTED" &&

        (!currentTimestamp || !j.time_end || differenceInSeconds(currentTimestamp, j.time_end) > jobLingerTime)

        (!currentTimestamp || !j.time_end || differenceInSeconds(currentTimestamp, j.time_end) < jobLingerTime)

      )

  ) as Job[];

  jobs = sortBy(jobs, j => j.state_current != "RUNNING", j => j.state_current, j => j.job_id);
@@ -234,7 +238,7 @@ export const useJobReplay = ({ 
      const nextQueryStart = addSeconds(queryStart, querySize)

      if (nextQueryStart < maxTimestamp!) {

        const nextQueryEnd = minDate([addSeconds(nextQueryStart, querySize), maxTimestamp!])

        queryClient.prefetchQuery({

        queryClient.prefetchInfiniteQuery({

          ...simulationSchedulerJobs(sim?.id ?? '', {

            start: nextQueryStart?.toISOString(), end: nextQueryEnd?.toISOString(),

            limit: 1000,
@@ -246,5 +250,14 @@ export const useJobReplay = ({ 
    }

  })



  return { data: jobs, maxTimestamp, currentTimestamp, nextTimestamp }

  return {

    data: jobs,

    maxTimestamp, currentTimestamp, nextTimestamp,

    hasNextPage,

    fetchNextPage: async () => {

      if (!isFetching) {

        await fetchNextPage();

      }

    },

  }

}
 
 No newline at end of file

src/util/queryOptions.ts

+10 −5
Original line number Diff line number Diff line 
import { keepPreviousData, queryOptions } from "@tanstack/react-query";

import { keepPreviousData, queryOptions, infiniteQueryOptions } from "@tanstack/react-query";

import axios from "axios";

import { Simulation } from "../models/Simulation.model";

import { CoolingCDU } from "../models/CoolingCDU.model";
@@ -96,13 +96,14 @@ export const simulationSystemStatsQueryOptions = ( 
 * Query jobs from the job endpoint.

 * sort should be passed as a list like ["desc:name", "asc:job_id"]

 * filters is in the format ["job_id=eq:5"]

 * This should be passed to useInfiniteQuery

 * TODO: allow passing sort/filters in cleaner format than raw query params.

 */

export const simulationSchedulerJobs = (

  simulationId: string,

  params?: {

    start?: string, end?: string,

    limit?: number, offset?: number,

    limit?: number,

    fields?: string[], sort?: string[],

    filters?: string[],

    currentTimestamp?: string,
@@ -113,21 +114,25 @@ export const simulationSchedulerJobs = ( 
  searchParams.push(...(params?.sort?.map(s => `sort=${s}`) ?? []));

  // Just specify filters as an array of ["job_id=eq:1"] for now

  searchParams.push(...(params?.filters ?? []));

  const limit = params?.limit ?? 1000;



  return queryOptions({

  return infiniteQueryOptions({

    queryKey: ["simulation", "scheduler", "jobs", simulationId, params],

    queryFn: async () => {

    queryFn: async ({ pageParam }) => {

      const res = await axios.get<{

        offset: number; limit: number, total_results: number, results: Job[];

      }>(`/frontier/simulation/${simulationId}/scheduler/jobs?${searchParams.join("&")}`, {

        params: {

          start: params?.start, end: params?.end,

          limit: params?.limit, offset: params?.offset,

          limit: limit, offset: pageParam,

          fields: params?.fields?.join(","),

        },

      });

      return res.data;

    },

    initialPageParam: 0,

    getNextPageParam: (lastPage, _allPages, lastPageParam, _allPageParams) =>

      (lastPageParam + limit < lastPage.total_results) ? lastPageParam + limit : undefined,

  })

}