docstr: Determine the number of tokens and KV blocks to reserve for a given request. Given a UID
  (this UID may not be recognized by the model yet), this will return the number of tokens
  and blocks to reserve for the request.
  
  Arguments:
      uid (int): The UID of the sequence (as tracked by the scheduling entity). If
          this is a new sequence (with a UID unknown to the inference engine), then
          an empty placeholder is created to pass to the occupancy logic.
      n_tokens (int): The number of tokens to hypothetically send.
  
  Returns:
      Tuple[int, Optional[int]]: Tuple of free kv blocks and the number of blocks
          required to schedule the sequence.
docstr_tok: ["Determine","the","number","of","tokens","and","KV","blocks","to","reserve","for","a","given","request",".","Given","a","UID","(","this","UID","may","not","be","recognized","by","the","model","yet",")",",","this","will","return","the","number","of","tokens","and","blocks","to","reserve","for","the","request",".","Arguments",":","uid","(","int",")",":","The","UID","of","the","sequence","(","as","tracked","by","the","scheduling","entity",")",".","If","this","is","a","new","sequence","(","with","a","UID","unknown","to","the","inference","engine",")",",","then","an","empty","placeholder","is","created","to","pass","to","the","occupancy","logic",".","n_tokens","(","int",")",":","The","number","of","tokens","to","hypothetically","send",".","Returns",":","Tuple","[","int",",","Optional","[","int","]","]",":","Tuple","of","free","kv","blocks","and","the","number","of","blocks","required","to","schedule","the","sequence","."]
code: def query(
      self, uid: int, max_request_tokens: int, max_request_blocks
  ) -> Tuple[int, torch.Tensor]:
      """
      Determine the number of tokens and KV blocks to reserve for a given request. Given a UID
      (this UID may not be recognized by the model yet), this will return the number of tokens
      and blocks to reserve for the request.
  
      Arguments:
          uid (int): The UID of the sequence (as tracked by the scheduling entity). If
              this is a new sequence (with a UID unknown to the inference engine), then
              an empty placeholder is created to pass to the occupancy logic.
          n_tokens (int): The number of tokens to hypothetically send.
  
      Returns:
          Tuple[int, Optional[int]]: Tuple of free kv blocks and the number of blocks
              required to schedule the sequence.
      """
      seq_desc = self._state_manager.get_sequence(uid)
      if seq_desc is None:
          if (
              self._state_manager.n_tracked_sequences
              == self._config.state_manager.max_tracked_sequences
          ):
              return (0, 0)
          seq_desc = PlaceholderSequenceDescriptor()
  
      req_tokens, req_blocks = self._model.get_kv_requirements(
          seq_desc, max_request_tokens, max_request_blocks
      )
  
      return (req_tokens, req_blocks)
  
imports: import os
  import json
  import pickle
  from typing import Iterable, Tuple
  import torch
  import deepspeed.comm
  from deepspeed.accelerator import get_accelerator
  from deepspeed.comm.comm import init_distributed
  from .model_implementations import InferenceV2Policy
  from .logging import inference_logger
  from .ragged import DSStateManager, RaggedBatchWrapper, PlaceholderSequenceDescriptor
  from .scheduling_utils import SchedulingError, SchedulingResult
  from .model_implementations.flat_model_helpers import make_param_filename, make_metadata_filename
  from .model_implementations.inference_model_base import DSInferenceModelBase
  from .config_v2 import RaggedInferenceEngineConfig
label_desc: Use image node_id 7 for calling the InferenceEngineV2 obj's underlying member method code with example usage: obj.query(uid, max_request_tokens, max_request_blocks) and returns: req_tokens,  req_blocks, int,  int