From 6197b260b8512379d5bb177d21566bad12cbe9e3 Mon Sep 17 00:00:00 2001
From: Richard Liaw <rliaw@berkeley.edu>
Date: Thu, 9 Nov 2017 00:56:20 -0800
Subject: [PATCH] Fix Jenkins issue introduced by Variant Generator (#1194)

* try fix

* shorten

* added a flag

* finish

* Fix linting.
---
 python/ray/tune/examples/tune_mnist_ray.py | 9 +++++++++
 test/jenkins_tests/run_multi_node_tests.sh | 6 ++----
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/python/ray/tune/examples/tune_mnist_ray.py b/python/ray/tune/examples/tune_mnist_ray.py
index 13f456783..0aeb5706f 100755
--- a/python/ray/tune/examples/tune_mnist_ray.py
+++ b/python/ray/tune/examples/tune_mnist_ray.py
@@ -215,6 +215,15 @@ if __name__ == '__main__':
         },
     }
 
+    # These arguments are only for testing purposes.
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--fast', action='store_true',
+                        help='Run minimal iterations.')
+    args, _ = parser.parse_known_args()
+
+    if args.fast:
+        spec['stop']['training_iteration'] = 2
+
     for trial in generate_trials(spec):
         runner.add_trial(trial)
 
diff --git a/test/jenkins_tests/run_multi_node_tests.sh b/test/jenkins_tests/run_multi_node_tests.sh
index 909898f27..97d3e5d30 100755
--- a/test/jenkins_tests/run_multi_node_tests.sh
+++ b/test/jenkins_tests/run_multi_node_tests.sh
@@ -129,7 +129,5 @@ docker run --shm-size=10G --memory=10G $DOCKER_SHA \
     python /ray/python/ray/rllib/test/test_checkpoint_restore.py
 
 docker run --shm-size=10G --memory=10G $DOCKER_SHA \
-    python /ray/python/ray/tune/tune.py \
-    --env mnist \
-    --stop '{"training_iteration": 2}' \
-    --config '{"script_file_path": "/ray/python/ray/tune/examples/tune_mnist_ray.py", "activation": "relu"}'
+    python /ray/python/ray/tune/examples/tune_mnist_ray.py \
+    --fast