class_docstr: Transformers handler class for testing
class_docstr_tok: ["Transformers","handler","class","for","testing"]
docstr: Expected behaviour: the sharded Bloom/Llama model is loaded.
  
  Args:
      ctx (context): It is a JSON Object containing information
      pertaining to the model artefacts parameters.
docstr_tok: ["Expected","behaviour",":","the","sharded","Bloom\/Llama","model","is","loaded",".","Args",":","ctx","(","context",")",":","It","is","a","JSON","Object","containing","information","pertaining","to","the","model","artefacts","parameters","."]
code: def initialize(self, ctx):
      """Expected behaviour: the sharded Bloom/Llama model is loaded.
  
      Args:
          ctx (context): It is a JSON Object containing information
          pertaining to the model artefacts parameters.
      """
      if ctx is not None or not hasattr(ctx, "model_yaml_config"):
          logger.error(
              "Context ctx and model-config are not appropriately passed in."
          )
  
      self.manifest = ctx.manifest
      gpu_id = ctx.system_properties.get("gpu_id", -1)
      model_dir = ctx.system_properties.get("model_dir")
  
      # Inference configs are collected together in model yaml config for handler use
      inference_config = ctx.model_yaml_config["handler"]
      self.inference_config = inference_config
      logger.info(self.inference_config)
  
      self.tp_size = self.inference_config.get("tp_size", 1)
      self.max_batch_size = self.inference_config.get(
          "max_batch_size", 4
      )
      self.max_input_len = self.inference_config.get(
          "max_input_len", 1024
      )
      self.max_output_len = self.inference_config.get(
          "max_output_len", 128
      )
  
      self.device = torch.device(
          "cuda:" + str(gpu_id)
          if torch.cuda.is_available() and gpu_id >= 0
          else "cpu"
      )
      logger.info(f"Device set to {self.device}")
      logger.info(
          f"torch.cuda.device_count() {torch.cuda.device_count()}"
      )
  
      # Unpacking from model_dir
      model_dir_path = os.path.join(model_dir, "model")
      with zipfile.ZipFile(model_dir + "/model.zip", "r") as zip_ref:
          zip_ref.extractall(model_dir_path)
      logger.info(
          f"Loading {self.inference_config['model_type']} pretrain model and tokenizer"
      )
      if self.inference_config["model_type"] == "bloom":
          self.model = BloomForCausalLM.from_pretrained(
              model_dir_path,
          )
          self.tokenizer = BloomTokenizerFast.from_pretrained(
              model_dir_path, return_tensors="pt"
          )
      elif self.inference_config["model_type"] == "llama":
          self.model = LlamaForCausalLM.from_pretrained(
              model_dir_path,
          )
          self.tokenizer = AutoTokenizer.from_pretrained(
              model_dir_path, return_tensors="pt"
          )
      else:
          logger.warning(
              f"Model type {self.inference_config['model_type']} not supported yet."
          )
  
      logger.info(
          "Transformer model from path %s loaded successfully",
          model_dir,
      )
  
      # NOTE world_size, rank, host, port here are used to launch colossalai dist environment
      # This world_size is different from the world size of TorchServe
      world_size = int(os.getenv("WORLD_SIZE", self.tp_size))
      assert (
          world_size == 1
      ), "Colossal-Inference with tensor parallel is not supported on TorchServe for now"
      rank = int(os.getenv("RANK", gpu_id))
      local_rank = int(os.getenv("LOCAL_RANK", gpu_id))
      host = os.getenv("MASTER_ADDR", "localhost")
      port = os.getenv(
          "MASTER_PORT", free_port()
      )  # use a random free port
  
      logger.info(
          f"  world_size {world_size}"
          f"  local_rank {local_rank}"
          f"  rank {rank}"
          f"  host {host}"
          f"  port {port}"
      )
  
      torch.cuda.set_device(self.device)
      self.model.half()
      self.model.cuda()
      self.model.eval()
  
      colossalai.launch(
          config={},
          rank=rank,
          world_size=world_size,
          host=host,
          port=port,
          backend="nccl",
      )
      logger.info("Initializing TPInferEngine ...")
      shard_config = ShardConfig(
          enable_tensor_parallelism=True if self.tp_size > 1 else False,
          extra_kwargs={"inference_only": True},
      )
      self.infer_engine = TPInferEngine(
          self.model,
          shard_config,
          self.max_batch_size,
          self.max_input_len,
          self.max_output_len,
      )
      logger.info("TPInferEngine initialized successfully")
  
      self.model = self.infer_engine.model
      self.initialized = True
  
file_path: ColossalAI/colossalai/legacy/inference/serving/torch_serve/Colossal_Inference_Handler.py
imports: import logging
  import os
  import zipfile
  from abc import ABC
  import torch
  import transformers
  from transformers import AutoTokenizer, BloomForCausalLM, BloomTokenizerFast, LlamaForCausalLM
  from ts.torch_handler.base_handler import BaseHandler
  import colossalai
  from colossalai.inference.tensor_parallel.engine import TPInferEngine
  from colossalai.shardformer import ShardConfig
  from colossalai.testing import free_port
label_desc: Use image node_id 2 for calling the ColossalInferenceHandler obj's underlying member method code with example usage: obj.initialize(ctx) without return types