enable running example apps in cluster mode (#357)

examples/alexnet/driver.py

@@ -7,15 +7,23 @@ import boto3
 import alexnet
 
 # Arguments to specify where the imagenet data is stored.
-parser = argparse.ArgumentParser(description="Parse information for data loading.")
+parser = argparse.ArgumentParser(description="Run the AlexNet example.")
+parser.add_argument("--node-ip-address", default=None, type=str, help="The IP address of this node.")
+parser.add_argument("--scheduler-address", default=None, type=str, help="The address of the scheduler.")
 parser.add_argument("--s3-bucket", required=True, type=str, help="Name of the bucket that contains the image data.")
 parser.add_argument("--key-prefix", default="ILSVRC2012_img_train/n015", type=str, help="Prefix for files to fetch.")
-parser.add_argument("--label-file", default="train.txt", type=str, help="File containing labels")
+parser.add_argument("--label-file", default="train.txt", type=str, help="File containing labels.")
 
 if __name__ == "__main__":
   args = parser.parse_args()
-  num_workers = 4
-  ray.init(start_ray_local=True, num_workers=num_workers)
+
+  # If node_ip_address and scheduler_address are provided, then this command
+  # will connect the driver to the existing scheduler. If not, it will start
+  # a local scheduler and connect to it.
+  ray.init(start_ray_local=(args.node_ip_address is None),
+           node_ip_address=args.node_ip_address,
+           scheduler_address=args.scheduler_address,
+           num_workers=(10 if args.node_ip_address is None else None))
 
   # Note we do not do sess.run(tf.initialize_all_variables()) because that would
   # result in a different initialization on each worker. Instead, we initialize
@@ -38,7 +46,7 @@ if __name__ == "__main__":
   filename_label_pairs = [line.split(" ") for line in filename_label_str]
   filename_label_dict = dict([(os.path.basename(name), label) for name, label in filename_label_pairs])
   filename_label_dict_id = ray.put(filename_label_dict)
-  print "Labels extracted"
+  print "Labels extracted."
 
   # Download the imagenet dataset.
   imagenet_data = alexnet.load_tarfiles_from_s3(args.s3_bucket, image_tar_files, [256, 256])
@@ -75,7 +83,8 @@ if __name__ == "__main__":
 
     # Launch tasks in parallel to compute the gradients for some batches.
     gradient_ids = []
-    for i in range(num_workers - 1):
+    num_batches = 4
+    for i in range(num_batches):
       # Choose a random batch and use it to compute the gradient of the loss.
       x_id, y_id = batches[np.random.randint(len(batches))]
       gradient_ids.append(alexnet.compute_grad.remote(x_id, y_id, mean_id, weights_id))

examples/hyperopt/README.md

@@ -64,7 +64,7 @@ def generate_random_params():
 results = []
 for _ in range(100):
   randparams = generate_random_params()
-  results.append((randparams, train_cnn_and_compute_accuracy(randparams, epochs)))
+  results.append((randparams, train_cnn_and_compute_accuracy(randparams, train_images, train_labels, validation_images, validation_labels)))
 ```
 
 Then we can inspect the contents of `results` and see which set of
@@ -105,7 +105,7 @@ computation. Instead, it simply submits a number of tasks to the scheduler.
 result_ids = []
 for _ in range(100):
   params = generate_random_params()
-  results.append((params, train_cnn_and_compute_accuracy.remote(params, epochs)))
+  results.append((params, train_cnn_and_compute_accuracy.remote(params, train_images, train_labels, validation_images, validation_labels)))
 ```
 
 If we wish to wait until the results have all been retrieved, we can retrieve

examples/hyperopt/driver.py

@@ -2,20 +2,34 @@
 # https://www.tensorflow.org/versions/r0.9/tutorials/mnist/pros/index.html#build-a-multilayer-convolutional-network
 import numpy as np
 import ray
-import os
+import argparse
 
 import tensorflow as tf
 from tensorflow.examples.tutorials.mnist import input_data
 
 import hyperopt
 
+parser = argparse.ArgumentParser(description="Run the hyperparameter optimization example.")
+parser.add_argument("--node-ip-address", default=None, type=str, help="The IP address of this node.")
+parser.add_argument("--scheduler-address", default=None, type=str, help="The address of the scheduler.")
+parser.add_argument("--trials", default=2, type=int, help="The number of random trials to do.")
+parser.add_argument("--steps", default=10, type=int, help="The number of steps of training to do per network.")
+
 if __name__ == "__main__":
-  ray.init(start_ray_local=True, num_workers=3)
+  args = parser.parse_args()
+
+  # If node_ip_address and scheduler_address are provided, then this command
+  # will connect the driver to the existing scheduler. If not, it will start
+  # a local scheduler and connect to it.
+  ray.init(start_ray_local=(args.node_ip_address is None),
+           node_ip_address=args.node_ip_address,
+           scheduler_address=args.scheduler_address,
+           num_workers=(10 if args.node_ip_address is None else None))
 
   # The number of sets of random hyperparameters to try.
-  trials = 2
+  trials = args.trials
   # The number of training passes over the dataset to use for network.
-  epochs = 10
+  steps = args.steps
 
   # Load the mnist data and turn the data into remote objects.
   print "Downloading the MNIST dataset. This may take a minute."
@@ -37,7 +51,7 @@ if __name__ == "__main__":
     dropout = np.random.uniform(0, 1)
     stddev = 10 ** np.random.uniform(-5, 5)
     params = {"learning_rate": learning_rate, "batch_size": batch_size, "dropout": dropout, "stddev": stddev}
-    results.append((params, hyperopt.train_cnn_and_compute_accuracy.remote(params, epochs, train_images, train_labels, validation_images, validation_labels)))
+    results.append((params, hyperopt.train_cnn_and_compute_accuracy.remote(params, steps, train_images, train_labels, validation_images, validation_labels)))
 
   # Fetch the results of the tasks and print the results.
   for i in range(trials):

examples/hyperopt/hyperopt.py

@@ -52,7 +52,7 @@ def cnn_setup(x, y, keep_prob, lr, stddev):
 # Define a remote function that takes a set of hyperparameters as well as the
 # data, consructs and trains a network, and returns the validation accuracy.
 @ray.remote([dict, int, np.ndarray, np.ndarray, np.ndarray, np.ndarray], [float])
-def train_cnn_and_compute_accuracy(params, epochs, train_images, train_labels, validation_images, validation_labels):
+def train_cnn_and_compute_accuracy(params, steps, train_images, train_labels, validation_images, validation_labels):
   # Extract the hyperparameters from the params dictionary.
   learning_rate = params["learning_rate"]
   batch_size = params["batch_size"]
@@ -68,7 +68,7 @@ def train_cnn_and_compute_accuracy(params, epochs, train_images, train_labels, v
   with tf.Session() as sess:
     # Initialize the network weights.
     sess.run(tf.initialize_all_variables())
-    for i in range(1, epochs):
+    for i in range(1, steps + 1):
       # Fetch the next batch of data.
       image_batch = get_batch(train_images, i, batch_size)
       label_batch = get_batch(train_labels, i, batch_size)
@@ -82,7 +82,7 @@ def train_cnn_and_compute_accuracy(params, epochs, train_images, train_labels, v
         if train_ac < 0.25:
           # Compute the validation accuracy and return.
           totalacc = accuracy.eval(feed_dict={x: validation_images, y: validation_labels, keep_prob: 1.0})
-          return totalacc
+          return float(totalacc)
     # Training is done, compute the validation accuracy and return.
     totalacc = accuracy.eval(feed_dict={x: validation_images, y: validation_labels, keep_prob: 1.0})
   return float(totalacc)

examples/lbfgs/driver.py

@@ -1,5 +1,5 @@
-import os
 import ray
+import argparse
 
 import numpy as np
 import scipy.optimize
@@ -7,8 +7,20 @@ import tensorflow as tf
 
 from tensorflow.examples.tutorials.mnist import input_data
 
+parser = argparse.ArgumentParser(description="Run the L-BFGS example.")
+parser.add_argument("--node-ip-address", default=None, type=str, help="The IP address of this node.")
+parser.add_argument("--scheduler-address", default=None, type=str, help="The address of the scheduler.")
+
 if __name__ == "__main__":
-  ray.init(start_ray_local=True, num_workers=16)
+  args = parser.parse_args()
+
+  # If node_ip_address and scheduler_address are provided, then this command
+  # will connect the driver to the existing scheduler. If not, it will start
+  # a local scheduler and connect to it.
+  ray.init(start_ray_local=(args.node_ip_address is None),
+           node_ip_address=args.node_ip_address,
+           scheduler_address=args.scheduler_address,
+           num_workers=(10 if args.node_ip_address is None else None))
 
   # Define the dimensions of the data and of the model.
   image_dimension = 784

examples/rl_pong/driver.py

@@ -4,9 +4,14 @@
 import numpy as np
 import cPickle as pickle
 import ray
+import argparse
 
 import gym
 
+parser = argparse.ArgumentParser(description="Run the Pong example.")
+parser.add_argument("--node-ip-address", default=None, type=str, help="The IP address of this node.")
+parser.add_argument("--scheduler-address", default=None, type=str, help="The address of the scheduler.")
+
 # hyperparameters
 H = 200 # number of hidden layer neurons
 batch_size = 10 # every how many episodes to do a param update?
@@ -108,7 +113,15 @@ def compute_gradient(model):
   return policy_backward(eph, epx, epdlogp, model), reward_sum
 
 if __name__ == "__main__":
-  ray.init(start_ray_local=True, num_workers=10)
+  args = parser.parse_args()
+
+  # If node_ip_address and scheduler_address are provided, then this command
+  # will connect the driver to the existing scheduler. If not, it will start
+  # a local scheduler and connect to it.
+  ray.init(start_ray_local=(args.node_ip_address is None),
+           node_ip_address=args.node_ip_address,
+           scheduler_address=args.scheduler_address,
+           num_workers=(10 if args.node_ip_address is None else None))
 
   # Run the reinforcement learning
   running_reward = None