code: def test_pipe_inference_schedule_firststage(micro_batches, stages=3):
      sched = schedule.InferenceSchedule(
          micro_batches=micro_batches, stages=stages, stage_id=0
      )
      assert sched.num_micro_batches == micro_batches
      full = list(iter(sched))
      for idx, cmds in enumerate(full):
          # Ensure we don't send an activation the first step
          if idx == 0:
              assert len(cmds) == 2
              assert type(cmds[0]) == schedule.LoadMicroBatch
              assert type(cmds[1]) == schedule.ForwardPass
              assert cmds[0].buffer_id == cmds[1].buffer_id
              continue
  
          # the last active step is only a send
          if idx == sched.num_micro_batches:
              assert len(cmds) == 1
              assert type(cmds[0]) == schedule.SendActivation
              continue
  
          # no work later on
          if idx > sched.num_micro_batches:
              assert len(cmds) == 0
              continue
  
          # Normally we need to load/forward/send
          assert len(cmds) == 3
          assert _count_type(cmds, schedule.LoadMicroBatch) == 1
          assert _count_type(cmds, schedule.ForwardPass) == 1
          assert _count_type(cmds, schedule.SendActivation) == 1
      assert len(full) == micro_batches + stages - 1
  
file_path: deepspeed/tests/unit/runtime/pipe/test_pipe_schedule.py
imports: import pytest
  import deepspeed.runtime.pipe.schedule
label_desc: Use image node_id 4 for calling a global function with example usage: test_pipe_inference_schedule_firststage(micro_batches, stages) without return types