code: def test_decoding_stage(monkeypatch):
      monkeypatch.syspath_prepend(
          (CURR_FILE_PATH / "test_data" / "streaming")
      )
  
      handler = StreamingHandler()
      ctx = MockContext(
          model_pt_file=None,
          model_dir=(
              CURR_FILE_PATH / "test_data" / "streaming"
          ).as_posix(),
          model_file="fake_streaming_model.py",
      )
      ctx.model_yaml_config["handler"] = {"modelId": "gpt2"}
  
      torch.manual_seed(42 * 42)
      handler.initialize(ctx)
  
      handler.context = ctx
  
      device = next(iter(handler.model.parameters())).device
  
      ctx.cache = {
          "id1": {
              "encoded": {
                  "input_ids": torch.randint(42, (1, 5), device=device),
                  "attention_mask": torch.ones(
                      (1, 5), dtype=int, device=device
                  ),
                  "past_key_values": None,
              },
          },
          "id2": {
              "encoded": {
                  "input_ids": torch.randint(42, (1, 8), device=device),
                  "attention_mask": torch.ones(
                      (1, 8), dtype=int, device=device
                  ),
                  "past_key_values": None,
              }
          },
      }
      ctx.cache["id1"]["encoded"]["attention_mask"][0, :2] = 0
  
      res = handler._run_prefill("id1")
      res = handler._run_prefill("id2")
  
      res = handler._run_decode(["id1"])
  
      assert len(res["id1"]["ids"]) == 1
      assert type(res["id1"]["text"]) == str
  
      assert ctx.cache["id1"]["encoded"]["input_ids"].size()[-1] == 5
      assert (
          ctx.cache["id1"]["encoded"]["attention_mask"].size()[-1] == 5
      )
  
      res = handler._run_decode(["id1", "id2"])
      assert ctx.cache["id1"]["encoded"]["input_ids"].size()[-1] == 10
      assert (
          ctx.cache["id1"]["encoded"]["attention_mask"].size()[-1] == 10
      )
  
      assert ctx.cache["id2"]["encoded"]["input_ids"].size()[-1] == 10
      assert (
          ctx.cache["id2"]["encoded"]["attention_mask"].size()[-1] == 10
      )
  
      res = handler._run_decode(["id1"])
      assert ctx.cache["id1"]["encoded"]["input_ids"].size()[-1] == 7
      assert (
          ctx.cache["id1"]["encoded"]["attention_mask"].size()[-1] == 7
      )
  
      res = handler._run_decode(["id1", "id2"])
      assert ctx.cache["id1"]["encoded"]["input_ids"].size()[-1] == 11
      assert (
          ctx.cache["id1"]["encoded"]["attention_mask"].size()[-1] == 11
      )
  
      assert ctx.cache["id2"]["encoded"]["input_ids"].size()[-1] == 11
      assert (
          ctx.cache["id2"]["encoded"]["attention_mask"].size()[-1] == 11
      )
  
file_path: serve/test/pytest/test_continuous_batching.py
imports: import json
  import shutil
  from pathlib import Path
  from queue import Empty
  from unittest.mock import patch
  from zipfile import ZIP_STORED, ZipFile
  import pytest
  import requests
  import test_utils
  import torch
  from model_archiver import ModelArchiverConfig
  from test_data.streaming.stream_handler import StreamingHandler
  from ts.torch_handler.unit_tests.test_utils.mock_context import MockContext
label_desc: Use image node_id 6 for calling a global function with example usage: test_decoding_stage(monkeypatch) without return types