[docker] Detect CPUs in container correctly (#10507)

Co-authored-by: simon-mo <simon.mo@hey.com>
Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
Co-authored-by: Alex Wu <itswu.alex@gmail.com>
This commit is contained in:
Ian Rodney
2020-09-13 23:40:48 -07:00
committed by GitHub
parent 660aee6311
commit 5bc2ba38fd
35 changed files with 173 additions and 45 deletions
+7 -7
View File
@@ -67,7 +67,7 @@ py_test(
srcs = ["tf/examples/tensorflow_train_example.py"],
tags = ["exclusive", "tf"],
deps = [":sgd_lib"],
args = ["--num-replicas=1"]
args = ["--num-replicas=1", "--smoke-test"]
)
py_test(
@@ -77,7 +77,7 @@ py_test(
srcs = ["tf/examples/tensorflow_train_example.py"],
tags = ["exclusive", "tf"],
deps = [":sgd_lib"],
args = ["--num-replicas=2"]
args = ["--num-replicas=2", "--smoke-test"]
)
py_test(
@@ -87,7 +87,7 @@ py_test(
srcs = ["tf/examples/tensorflow_train_example.py"],
tags = ["exclusive", "tf"],
deps = [":sgd_lib"],
args = ["--tune"]
args = ["--tune", "--smoke-test"]
)
# --------------------------------------------------------------------
@@ -148,7 +148,7 @@ py_test(
srcs = ["torch/examples/train_example.py"],
tags = ["exclusive", "pytorch"],
deps = [":sgd_lib"],
args = ["--num-workers=1"]
args = ["--num-workers=1", "--smoke-test"]
)
py_test(
@@ -158,7 +158,7 @@ py_test(
srcs = ["torch/examples/train_example.py"],
tags = ["exclusive", "pytorch"],
deps = [":sgd_lib"],
args = ["--num-workers=2"]
args = ["--num-workers=2", "--smoke-test"]
)
py_test(
@@ -168,7 +168,7 @@ py_test(
srcs = ["torch/examples/tune_example.py"],
tags = ["exclusive", "pytorch"],
deps = [":sgd_lib"],
args = ["--num-workers=1"]
args = ["--num-workers=1", "--smoke-test"]
)
py_test(
@@ -178,7 +178,7 @@ py_test(
srcs = ["torch/examples/tune_example.py"],
tags = ["exclusive", "pytorch"],
deps = [":sgd_lib"],
args = ["--num-workers=2"]
args = ["--num-workers=2", "--smoke-test"]
)
@@ -177,7 +177,10 @@ if __name__ == "__main__":
help="Finish quickly for testing. Assume False for users.")
args, _ = parser.parse_known_args()
ray.init(address=args.address)
if args.smoke_test:
ray.init(num_cpus=2)
else:
ray.init(address=args.address)
data_size = 60000
test_size = 10000
batch_size = args.batch_size
@@ -114,6 +114,8 @@ def tune_example(num_replicas=1, use_gpu=False):
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--smoke-test", action="store_true", help="Finish quickly for testing")
parser.add_argument(
"--address",
required=False,
@@ -135,7 +137,10 @@ if __name__ == "__main__":
args, _ = parser.parse_known_args()
ray.init(address=args.address)
if args.smoke_test:
ray.init(num_cpus=2)
else:
ray.init(address=args.address)
if args.tune:
tune_example(num_replicas=args.num_replicas, use_gpu=args.use_gpu)
@@ -109,7 +109,10 @@ class Training(TrainingOperator):
if __name__ == "__main__":
ray.init(address=None if args.local else "auto")
if args.local:
ray.init(num_cpus=2)
else:
ray.init(address="auto")
num_workers = 2 if args.local else int(ray.cluster_resources().get(device))
from ray.util.sgd.torch.examples.train_example import LinearDataset
+4 -1
View File
@@ -289,7 +289,10 @@ if __name__ == "__main__":
default=False,
help="Enables GPU training")
args = parser.parse_args()
ray.init(address=args.address)
if args.smoke_test:
ray.init(num_cpus=2)
else:
ray.init(address=args.address)
trainer = train_example(
num_workers=args.num_workers,
@@ -128,8 +128,10 @@ def main():
setup_default_logging()
args, args_text = parse_args()
ray.init(address=args.ray_address)
if args.smoke_test:
ray.init(num_cpus=int(args.ray_num_workers))
else:
ray.init(address=args.ray_address)
CustomTrainingOperator = TrainingOperator.from_creators(
model_creator=model_creator,
@@ -115,10 +115,17 @@ if __name__ == "__main__":
help="Enables GPU training")
parser.add_argument(
"--tune", action="store_true", default=False, help="Tune training")
parser.add_argument(
"--smoke-test",
action="store_true",
default=False,
help="Finish quickly for testing.")
args, _ = parser.parse_known_args()
import ray
ray.init(address=args.address)
if args.smoke_test:
ray.init(num_cpus=2)
else:
ray.init(address=args.address)
train_example(num_workers=args.num_workers, use_gpu=args.use_gpu)
@@ -59,6 +59,8 @@ def tune_example(operator_cls, num_workers=1, use_gpu=False):
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument(
"--smoke-test", action="store_true", help="Finish quickly for testing")
parser.add_argument(
"--address",
type=str,
@@ -77,7 +79,10 @@ if __name__ == "__main__":
args, _ = parser.parse_known_args()
ray.init(address=args.address)
if args.smoke_test:
ray.init(num_cpus=2)
else:
ray.init(address=args.address)
CustomTrainingOperator = TrainingOperator.from_creators(
model_creator=model_creator, optimizer_creator=optimizer_creator,
data_creator=data_creator, loss_creator=nn.MSELoss)