old_contents: # Lint as: python3
  # Copyright 2020 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """TF2 parameter server training utilities.
  
  Parameter server training in TF2 is currently under development.
  """
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl import logging
  from tensorflow.python.training import server_lib
  
  
  def start_server(cluster_resolver, protocol):
    """Start a server and block the process from exiting."""
    # This function is for multi-processing test or users who would like to have
    # every job run the same binary for simplicity.
    assert (cluster_resolver.task_type == 'worker' or
            cluster_resolver.task_type == 'ps')
    server = server_lib.Server(
        cluster_resolver.cluster_spec().as_cluster_def(),
        job_name=cluster_resolver.task_type,
        task_index=cluster_resolver.task_id,
        protocol=protocol)
  
    logging.info('TensorFlow server started for job %s, task %d.',
                 cluster_resolver.task_type, cluster_resolver.task_id)
  
    # Blocking the process that starts a server from exiting.
    server.join()
  
new_contents: # Lint as: python3
  # Copyright 2020 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """TF2 parameter server training utilities.
  
  Parameter server training in TF2 is currently under development.
  """
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl import logging
  from tensorflow.python.training import server_lib
  
  
  def start_server(cluster_resolver, protocol):
    """Start a server and block the process from exiting."""
    # This function is for multi-processing test or users who would like to have
    # every job run the same binary for simplicity.
    if not (cluster_resolver.task_type == 'worker' or
            cluster_resolver.task_type == 'ps'):
      raise ValueError('Unexpected task_type to start a server: {}'.format(
          cluster_resolver.task_type))
  
    server = server_lib.Server(
        cluster_resolver.cluster_spec().as_cluster_def(),
        job_name=cluster_resolver.task_type,
        task_index=cluster_resolver.task_id,
        protocol=protocol)
  
    logging.info('TensorFlow server started for job %s, task %d.',
                 cluster_resolver.task_type, cluster_resolver.task_id)
  
    # Blocking the process that starts a server from exiting.
    server.join()
  
subject: Replace an assert with a more informative ValueError.