docstr: Put a ragged batch onto the inference engine. This will perform one forward and return
  a Tensor of the shape [len(batch_uids), *output_shape]. Logits for the non-final tokens
  are not calculated.
  
  Arguments:
      batch_uids: Iterable of uids for the batch on the host
      batch_tokens: Iterable of token tensors for the batch on the host
docstr_tok: ["Put","a","ragged","batch","onto","the","inference","engine",".","This","will","perform","one","forward","and","return","a","Tensor","of","the","shape","[","len","(","batch_uids",")",",","*","output_shape","]",".","Logits","for","the","non-final","tokens","are","not","calculated",".","Arguments",":","batch_uids",":","Iterable","of","uids","for","the","batch","on","the","host","batch_tokens",":","Iterable","of","token","tensors","for","the","batch","on","the","host"]
code: def put(
      self,
      batch_uids: Iterable[int],
      batch_tokens: Iterable[torch.Tensor],
  ) -> torch.Tensor:
      """
      Put a ragged batch onto the inference engine. This will perform one forward and return
      a Tensor of the shape [len(batch_uids), *output_shape]. Logits for the non-final tokens
      are not calculated.
  
      Arguments:
          batch_uids: Iterable of uids for the batch on the host
          batch_tokens: Iterable of token tensors for the batch on the host
      """
  
      token_lens = [len(tokens) for tokens in batch_tokens]
      schedule_check = self.can_schedule(batch_uids, token_lens)
      if schedule_check != SchedulingResult.Success:
          raise SchedulingError(schedule_check)
  
      self._batch.clear()
      for uid, tokens in zip(batch_uids, batch_tokens):
          host_seq_desc = self._state_manager.get_or_create_sequence(
              uid
          )
          self._model.maybe_allocate_kv(host_seq_desc, tokens.numel())
          host_seq_desc.pre_forward(tokens.numel())
  
          # We can disable checks since we already validated schedulability.
          self._batch.insert_sequence(
              host_seq_desc, tokens, do_checks=False
          )
  
      # Send all metadata to the device
      self._batch.finalize()
  
      # Prep all data structures for the actual forward (in anticipation of CG in the future)
      # and also to amortize some of the costs in a more straightforward way.
      self._model.prepare_batch(self._batch)
  
      # Model implementation will pick up in the forward.
      logits = self._model.forward(self._batch)
  
      # We return one set of logits per sequence in the batch (saves cost on unembedding)
      assert logits.shape[0] == self._batch.current_sequences
  
      for uid in batch_uids:
          host_seq_desc = self._state_manager.get_sequence(uid)
          host_seq_desc.post_forward()  # Updates sequence metadata.
          self._model.maybe_free_kv(host_seq_desc)
  
      return logits
  
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 6 for calling the InferenceEngineV2 obj's underlying member method code with example usage: obj.put(batch_uids, batch_tokens) and returns: logits