code: def test_preempt_workers_ask_chief(
      dummy: bool, auto_ack: bool
  ) -> None:
      with parallel.Execution(2) as pex:
  
          @pex.run
          def do_test() -> None:
              if not dummy:
                  state, context = make_test_preempt_context(
                      pex.distributed, core.PreemptMode.WorkersAskChief
                  )
              else:
                  context = core.DummyPreemptContext(
                      pex.distributed, core.PreemptMode.WorkersAskChief
                  )
  
              with context:
                  if pex.rank == 0:
                      # Check preemption.
                      assert context.should_preempt() is False
                      # Make sure the worker is receiving broadcasts.
                      _ = pex.distributed.broadcast(False)
                      if not dummy:
                          # No ack preemption calls yet.
                          state.mock_session.post.assert_not_called()
                          # Send the preemption signal.
                          state.preempt()
                          wait_on_watcher(context)
                          assert (
                              context.should_preempt(auto_ack=auto_ack)
                              is True
                          )
                          # Call again, to make sure we only ack once.
                          assert (
                              context.should_preempt(auto_ack=auto_ack)
                              is True
                          )
                          if auto_ack:
                              state.mock_session.post.assert_called_once()
                          else:
                              state.mock_session.post.assert_not_called()
                  else:
                      # Intercept the broadcast from the chief to make sure it's happening.
                      out = pex.distributed.broadcast(None)
                      assert out is False, out
                      # Try receving from the chief.
                      assert context.should_preempt() is False
                      if not dummy:
                          # The chief should send a True now.
                          assert context.should_preempt() is True
                          # Only the chief acknowledges the preemption signal.
                          state.mock_session.post.assert_not_called()
  
file_path: determined/harness/tests/core/test_preempt.py
imports: import threading
  import time
  from typing import Any, Dict, Tuple
  from unittest import mock
  import pytest
  from determined import core
  from tests import parallel
label_desc: Use image node_id 3 for calling a global function with example usage: test_preempt_workers_ask_chief(dummy, auto_ack) without return types