分布式TensorFlow问题：无法运行仅评估工作节点

Hey there! Let’s tackle this distributed TensorFlow setup— I’ve been down this road before, so I know how finicky cluster configs can get when using tf.contrib.learn.Experiment. The core issue here is usually mixing up the roles of training nodes, parameter servers, and evaluators, plus not setting up the checkpoint triggering correctly for the evaluator.

1. First, Clarify Your Cluster Roles #

Let’s lock in what each node should do to avoid confusion:

• PC1: Runs two processes:
  • Parameter Server (ps): Manages shared model parameters
  • Worker Node (worker): Runs GPU training, feeds updates to the PS
• PC2: Runs a single Evaluator process: Only runs evaluation when a new checkpoint is saved by PC1’s worker (no need to join the training cluster)

2. Correct Cluster Spec & Command Line Parsing #

First, update your main function to properly parse job roles and cluster details. The key here is that the evaluator doesn’t need to be part of the training cluster spec— it just needs access to the checkpoint directory.

Here’s a cleaned-up main function snippet:

import tensorflow as tf
from tensorflow.contrib.learn import Experiment

def main(unused_argv):
    # Parse command line arguments
    parser = tf.flags.ArgumentParser()
    parser.add_argument("--job_name", type=str, required=True, help="ps/worker/evaluator")
    parser.add_argument("--task_index", type=int, default=0, help="Task index for ps/worker")
    parser.add_argument("--cluster_spec", type=str, default='{"ps": ["pc1_ip:2222"], "worker": ["pc1_ip:2223"]}', help="JSON cluster spec")
    parser.add_argument("--checkpoint_dir", type=str, default="/path/to/shared/checkpoints", help="Checkpoint storage path")
    FLAGS = parser.parse_args()

    # Parse cluster spec from JSON
    cluster = tf.train.ClusterSpec(eval(FLAGS.cluster_spec))

    # Handle each job type
    if FLAGS.job_name == "ps":
        # Start parameter server
        server = tf.train.Server(cluster, job_name="ps", task_index=FLAGS.task_index)
        server.join()
    elif FLAGS.job_name == "worker":
        # Configure training worker (GPU on PC1)
        config = tf.contrib.learn.RunConfig(
            model_dir=FLAGS.checkpoint_dir,
            session_config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True),
            cluster_spec=cluster,
            task_type="worker",
            task_id=FLAGS.task_index
        )
        # Assume you have your model_fn and input_fns defined
        experiment = Experiment(
            estimator=your_model_fn(config=config),
            train_input_fn=train_input_fn,
            eval_input_fn=eval_input_fn,
            train_steps=10000,
            eval_steps=100,
            train_config=config
        )
        experiment.train()
    elif FLAGS.job_name == "evaluator":
        # Configure evaluator (PC2, no cluster membership)
        eval_config = tf.contrib.learn.RunConfig(
            model_dir=FLAGS.checkpoint_dir,
            session_config=tf.ConfigProto(allow_soft_placement=True)
        )
        experiment = Experiment(
            estimator=your_model_fn(config=eval_config),
            train_input_fn=None,  # No training needed
            eval_input_fn=eval_input_fn,
            eval_steps=100,
            eval_config=tf.contrib.learn.EvalConfig(
                wait_for_checkpoint=True,  # Wait for new checkpoints
                throttle_secs=30,  # Check for new checkpoints every 30s
                start_delay_secs=60  # Wait 1min before first eval
            )
        )
        experiment.continuous_eval()  # Run eval repeatedly on new checkpoints
    else:
        raise ValueError(f"Unknown job name: {FLAGS.job_name}")

if __name__ == "__main__":
    tf.app.run()

3. Critical Configuration Notes #

• GPU Training on PC1: Make sure the worker process uses the GPU by either setting CUDA_VISIBLE_DEVICES=0 in the environment, or adding with tf.device('/gpu:0'): blocks in your model function.
• Checkpoint Sharing: The checkpoint_dir must be accessible to both PC1 and PC2. Use a network share (like NFS), or sync checkpoints to a cloud storage bucket if you’re using cloud instances.
• Evaluator Trigger: The EvalConfig settings are make-or-break:
  • wait_for_checkpoint=True: Ensures the evaluator doesn’t exit immediately if no checkpoint exists yet.
  • throttle_secs: Controls how often it checks for new checkpoints (adjust based on your training speed).
  • start_delay_secs: Gives the training process time to save the first checkpoint.

4. Launch Commands for Each Process #

On PC1 (GPU Node): #

1. Start Parameter Server:

python your_script.py --job_name=ps --task_index=0 --cluster_spec='{"ps": ["pc1_ip:2222"], "worker": ["pc1_ip:2223"]}' --checkpoint_dir="/mnt/shared/checkpoints"

2. Start Training Worker (GPU):

CUDA_VISIBLE_DEVICES=0 python your_script.py --job_name=worker --task_index=0 --cluster_spec='{"ps": ["pc1_ip:2222"], "worker": ["pc1_ip:2223"]}' --checkpoint_dir="/mnt/shared/checkpoints"

On PC2 (Evaluation Node): #

python your_script.py --job_name=evaluator --checkpoint_dir="//pc1_ip/shared/checkpoints"

5. Fixing Common Issues #

• Cluster Connection Failures: Ensure PC1’s firewall allows incoming connections on ports 2222/2223, and that pc1_ip is reachable from PC2.
• Evaluator Doesn’t Run: Double-check that wait_for_checkpoint=True is set, and that the checkpoint directory is accessible. If the first checkpoint takes time to save, increase start_delay_secs.
• Worker Doesn’t Use GPU: Verify that TensorFlow is built with GPU support, and that CUDA_VISIBLE_DEVICES is set correctly. Check the logs for Device mapping lines to confirm GPU usage.
• PS/Worker Sync Issues: Make sure the cluster spec is identical across all training processes (ps and worker). A typo in the IP/port will break communication.

内容的提问来源于stack exchange，提问作者GPhilo