code: def test(
      self,
      model,
      task,
      query,
      inf_kwargs,
      cuda_graphs,
      use_cuda_events,
      dtype=torch.float16,
  ):
      if cuda_graphs and "bert" not in model:
          pytest.skip(f"CUDA Graph not supported for {model}")
  
      local_rank = int(os.getenv("LOCAL_RANK", "0"))
      world_size = int(os.getenv("WORLD_SIZE", "1"))
  
      pipe = pipeline(task, model, framework="pt", device=local_rank)
      pipe.model = deepspeed.init_inference(
          pipe.model,
          dtype=dtype,
          mp_size=world_size,
          replace_with_kernel_inject=True,
          replace_method="auto",
          enable_cuda_graph=cuda_graphs,
      )
      pipe.model.profile_model_time(use_cuda_events=use_cuda_events)
  
      e2e_times = []
      model_times = []
      for _ in range(10):
          torch.cuda.synchronize()
          start = time.perf_counter_ns()
  
          r = pipe(query, **inf_kwargs)
  
          torch.cuda.synchronize()
          end = time.perf_counter_ns()
  
          e2e_times.append((end - start) / 1e6)  # convert ns to ms
          model_times.extend(pipe.model.model_times())
  
      for e2e_t, model_t in zip(e2e_times, model_times):
          assert e2e_t >= model_t
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/inference/test_model_profiling.py
imports: import os
  import time
  import pytest
  import torch
  import deepspeed
  from transformers import pipeline
  from unit.common import DistributedTest
label_desc: Use image node_id 1 for calling the TestModelProfiling obj's underlying member method code with example usage: obj.test(model, task, query, inf_kwargs, cuda_graphs, use_cuda_events, dtype) without return types