code: def test_sweep_scheduler_base_run_states(
      user, relay_server, sweep_config, monkeypatch
  ):
      _patch_wandb_run(monkeypatch)
      with relay_server():
          _entity = user
          _project = "test-project"
          api = internal.Api()
          sweep_id = wandb.sweep(
              sweep_config, entity=_entity, project=_project
          )
  
          # Mock api.get_run_state() to return crashed and running runs
          mock_run_states: Dict[str, RunState] = {
              "run1": RunState.CRASHED,
              "run2": RunState.FAILED,
              "run3": RunState.KILLED,
              "run4": RunState.FINISHED,
              "run5": RunState.RUNNING,
              "run6": RunState.PENDING,
              "run7": RunState.PREEMPTED,
              "run8": RunState.PREEMPTING,
              "run9": RunState.UNKNOWN,
              "run10": "?????",
          }
  
          def mock_get_run_state(
              entity, project, run_id, *args, **kwargs
          ):
              if run_id not in mock_run_states:  # is scheduler
                  return RunState.RUNNING
              return mock_run_states[run_id]
  
          api.get_run_state = mock_get_run_state
          _scheduler = Scheduler(
              api, sweep_id=sweep_id, entity=_entity, project=_project
          )
          # Load up the runs into the Scheduler run dict
          for i, run_id in enumerate(mock_run_states.keys()):
              _scheduler._runs[run_id] = SweepRun(
                  id=run_id, state=RunState.RUNNING, worker_id=i
              )
          _scheduler._update_run_states()
          for run_id, _state in mock_run_states.items():
              if (
                  _state == "?????"
              ):  # unknown state should be considered alive, but unknown
                  assert (
                      _scheduler._runs[run_id].state == RunState.UNKNOWN
                  )
                  continue
              if not _state.is_alive:
                  # Dead runs should be removed from the run dict
                  assert run_id not in _scheduler._runs.keys()
              else:
                  assert _scheduler._runs[run_id].state == _state
  
          # ---- If get_run_state errors out, runs should have the state UNKNOWN
          def mock_get_run_state_raise_exception(*args, **kwargs):
              raise CommError("Generic Exception")
  
          api.get_run_state = mock_get_run_state_raise_exception
          sweep_id = wandb.sweep(
              sweep_config, entity=_entity, project=_project
          )
          _scheduler = Scheduler(
              api, sweep_id=sweep_id, entity=_entity, project=_project
          )
          _scheduler._runs["foo_run_1"] = SweepRun(
              id="foo_run_1", state=RunState.RUNNING, worker_id=1
          )
          _scheduler._runs["foo_run_2"] = SweepRun(
              id="foo_run_2", state=RunState.RUNNING, worker_id=2
          )
          _scheduler._update_run_states()
          assert _scheduler._runs["foo_run_1"].state == RunState.UNKNOWN
          assert _scheduler._runs["foo_run_2"].state == RunState.UNKNOWN
  
imports: import asyncio
  from typing import Dict
  from unittest.mock import Mock, patch
  import pytest
  import wandb
  from wandb.apis import internal, public
  from wandb.errors import CommError
  from wandb.sdk.launch.sweeps import SchedulerError, load_scheduler
  from wandb.sdk.launch.sweeps.scheduler import RunState, Scheduler, SchedulerState, SweepRun
  from wandb.sdk.launch.sweeps.scheduler_sweep import SweepScheduler
  from wandb.sdk.launch.sweeps.utils import construct_scheduler_args
  from .test_wandb_sweep import VALID_SWEEP_CONFIGS_MINIMAL
label_desc: Use image node_id 9 for calling a global function with example usage: test_sweep_scheduler_base_run_states(user, relay_server, sweep_config, monkeypatch) and returns: mock_run_states, RunState