code: def compute_schedule(pages_in_batch, prefetch_ahead=3):
      # We determine what is the early and latest times we will need a page
      page_end = {}
      page_start = {}
      for b_id, pages in enumerate(pages_in_batch):
          for page in pages:
              page_end[page] = b_id
              if page not in page_start:
                  page_start[page] = b_id
  
      # We determine which pages are
      # - Can be preloaded
      # - Are needed
      # - Can be diposed of
      # At a given batch
      entering_at = defaultdict(set)
      can_prefetch_at = defaultdict(set)
      leaving_at = defaultdict(set)
      for page in page_start.keys():
          prefetch_start = max(0, page_start[page] - prefetch_ahead)
          can_prefetch_at[prefetch_start].add(page)
          entering_at[page_start[page]].add(page)
          leaving_at[page_end[page] + 1].add(page)
  
      # We now find how many pages we need to keep in our buffer
      # We also determine where which page is going to reside
      next_slot = 0
      page_to_slot = {}
      free_slots = set()
  
      # For each batch
      for b_id in range(len(pages_in_batch)):
          # First we free the pages that are leaving
          for page in leaving_at[b_id]:
              free_slots.add(page_to_slot[page])
  
          # We use the prefetch timing here because we want to be able
          # To start prefetching ahead of time and not overwrite a slot
          # That is currently used
          for page in can_prefetch_at[b_id]:
              # Then we find a slot for the incoming pages
              if free_slots:
                  # There is a slot available for this page
                  slot = free_slots.pop()
              else:
                  # We have to allocate a new slot because we ran out
                  slot = next_slot
                  next_slot += 1
  
              page_to_slot[page] = slot
  
      return Schedule(
          next_slot,
          page_to_slot,
          can_prefetch_at,
          entering_at,
          leaving_at,
      )
  
file_path: ffcv/ffcv/memory_managers/process_cache/schedule.py
imports: from collections import defaultdict
  from dataclasses import dataclass
  from typing import Mapping
  from queue import Queue
  import numpy
  from .page_reader import PageReader
label_desc: Use image node_id 1 for calling a global function with example usage: compute_schedule(pages_in_batch, prefetch_ahead) and returns: Schedule