mirror of
https://github.com/wassname/ray.git
synced 2026-07-13 09:15:06 +08:00
[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:
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user