From 45d656075991df7691739c78ee78f35ca304290a Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Thu, 15 Apr 2021 16:10:29 +0200 Subject: [PATCH] [RLlib] Fix flakey custom_fast_model_torch/tf tests. (#15330) --- rllib/BUILD | 21 ++++++++++----------- rllib/examples/custom_fast_model.py | 2 +- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/rllib/BUILD b/rllib/BUILD index 7f96e6947..be832e0aa 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1889,20 +1889,19 @@ py_test( name = "examples/custom_fast_model_tf", main = "examples/custom_fast_model.py", tags = ["examples", "examples_C"], - size = "small", + size = "medium", srcs = ["examples/custom_fast_model.py"], - args = ["--stop-iters=1", "--num-cpus=4"] + args = ["--stop-iters=1"] ) -# Skip because it is consistently failing in the master. -# py_test( -# name = "examples/custom_fast_model_torch", -# main = "examples/custom_fast_model.py", -# tags = ["examples", "examples_C"], -# size = "small", -# srcs = ["examples/custom_fast_model.py"], -# args = ["--torch", "--stop-iters=1", "--num-cpus=4"] -# ) +py_test( + name = "examples/custom_fast_model_torch", + main = "examples/custom_fast_model.py", + tags = ["examples", "examples_C"], + size = "medium", + srcs = ["examples/custom_fast_model.py"], + args = ["--stop-iters=1", "--torch"] +) py_test( name = "examples/custom_keras_model_a2c", diff --git a/rllib/examples/custom_fast_model.py b/rllib/examples/custom_fast_model.py index e9e8a1f98..2df8653f2 100644 --- a/rllib/examples/custom_fast_model.py +++ b/rllib/examples/custom_fast_model.py @@ -15,7 +15,7 @@ from ray.rllib.examples.models.fast_model import FastModel, TorchFastModel from ray.rllib.models import ModelCatalog parser = argparse.ArgumentParser() -parser.add_argument("--num-cpus", type=int, default=2) +parser.add_argument("--num-cpus", type=int, default=4) parser.add_argument("--torch", action="store_true") parser.add_argument("--stop-iters", type=int, default=200) parser.add_argument("--stop-timesteps", type=int, default=100000)