code: def test_pipe_inference_schedule_midstage(micro_batches, stages=3):
      sched = schedule.InferenceSchedule(
          micro_batches=micro_batches, stages=stages, stage_id=1
      )
  
      full = list(iter(sched))
      for idx, cmds in enumerate(full):
          if idx < sched.stage:
              assert len(cmds) == 0
              continue
          if idx == sched.stage + sched.num_micro_batches:
              assert len(cmds) == 1
              assert type(cmds[0]) == schedule.SendActivation
              continue
          if idx > sched.stage + sched.num_micro_batches:
              assert len(cmds) == 0
              continue
          assert _count_type(cmds, schedule.LoadMicroBatch) == 0
          assert _count_type(cmds, schedule.ForwardPass) == 1
          assert _count_type(cmds, schedule.RecvActivation) == 1
          if idx > sched.stage:
              assert _count_type(cmds, schedule.SendActivation) == 1
      assert len(full) == micro_batches + stages - 1
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/runtime/pipe/test_pipe_schedule.py
imports: import pytest
  import deepspeed.runtime.pipe.schedule
label_desc: Use image node_id 5 for calling a global function with example usage: test_pipe_inference_schedule_midstage(micro_batches, stages) without return types