code: def test_preempt_chief_only(dummy: bool, auto_ack: bool) -> None:
      with parallel.Execution(2) as pex:
          # Steal the automatically-created pex.distributed contexts, then test chief/worker serially
          # so we know they're not using distributed comms.
          @pex.run
          def distributed_contexts() -> core.DistributedContext:
              return pex.distributed
  
          # Test chief.
          if not dummy:
              state, context = make_test_preempt_context(
                  distributed_contexts[0], core.PreemptMode.ChiefOnly
              )
          else:
              context = core.DummyPreemptContext(
                  distributed_contexts[0], core.PreemptMode.ChiefOnly
              )
          with context:
              assert context.should_preempt() is 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()
  
          # Test worker.
          if not dummy:
              state, context = make_test_preempt_context(
                  distributed_contexts[1], core.PreemptMode.ChiefOnly
              )
          else:
              context = core.DummyPreemptContext(
                  distributed_contexts[1], core.PreemptMode.ChiefOnly
              )
          with context:
              with pytest.raises(
                  RuntimeError,
                  match="should_preempt.*called from non-chief",
              ):
                  context.should_preempt()
  
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 4 for calling a global function with example usage: test_preempt_chief_only(dummy, auto_ack) and returns: pex