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