code: def test_pipe_stagequery():
      sched = schedule.TrainSchedule(
          stages=3, micro_batches=4, stage_id=0
      )
      assert sched.is_first_stage
      assert not sched.is_last_stage
  
      sched = schedule.TrainSchedule(
          stages=3, micro_batches=4, stage_id=1
      )
      assert not sched.is_first_stage
      assert not sched.is_last_stage
  
      sched = schedule.TrainSchedule(
          stages=3, micro_batches=4, stage_id=2
      )
      assert not sched.is_first_stage
      assert sched.is_last_stage
  
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 9 for calling a global function with example usage: test_pipe_stagequery() without return types