code: def _set_model_dir_for_run_config(model_dir=None):
      """ContextManager for overwriting environment configuration for RunConfig."""
      old_tf_config_str = os.environ.get(_TF_CONFIG_ENV)
  
      new_tf_config = (
          copy.deepcopy(json.loads(old_tf_config_str))
          if old_tf_config_str
          else {}
      )
  
      if model_dir is not None:
          new_tf_config[_MODEL_DIR_KEY] = model_dir
  
      if FLAGS.phoenix_master is not None:
          new_tf_config[_SESSION_MASTER_KEY] = FLAGS.phoenix_master
  
      os.environ[_TF_CONFIG_ENV] = json.dumps(new_tf_config)
      try:
          yield
  
      finally:
          if old_tf_config_str is not None:
              os.environ[_TF_CONFIG_ENV] = old_tf_config_str
          else:
              del os.environ[_TF_CONFIG_ENV]
  
imports: import contextlib
  import copy
  import json
  import os
  import random
  from absl import flags
  from absl import logging
  import keras_tuner
  from model_search import hparam
  from model_search import phoenix
  from model_search import registry
  from model_search.data import data
  from model_search.metadata import ml_metadata_db
  from model_search.proto import phoenix_spec_pb2
  import tensorflow.compat.v2
  from google.protobuf import text_format
label_desc: Use image node_id 1 for calling a global function with example usage: _set_model_dir_for_run_config(model_dir) without return types