def determine_context(args: argparse.Namespace,
exit_stack: ExitStack) -> List[mx.Context]:
"""
Determine the context we should run on (CPU or GPU).
:param args: Arguments as returned by argparse.
:param exit_stack: An ExitStack from contextlib.
:return: A list with the context(s) to run on.
"""
if args.use_cpu:
logger.info("Training Device: CPU")
context = [mx.cpu()]
else:
num_gpus = utils.get_num_gpus()
check_condition(
num_gpus >= 1,
"No GPUs found, consider running on the CPU with --use-cpu "
"(note: check depends on nvidia-smi and this could also mean that the nvidia-smi "
"binary isn't on the path).")
if args.disable_device_locking:
context = utils.expand_requested_device_ids(args.device_ids)
else:
context = exit_stack.enter_context(
utils.acquire_gpus(args.device_ids, lock_dir=args.lock_dir))
if args.batch_type == C.BATCH_TYPE_SENTENCE:
check_condition(
args.batch_size % len(context) == 0,
"When using multiple devices the batch size must be "
"divisible by the number of devices. Choose a batch "
"size that is a multiple of %d." % len(context))
logger.info("Training Device(s): GPU %s", context)
context = [mx.gpu(gpu_id) for gpu_id in context]
return context
def main():
logging.basicConfig(level=logging.INFO)
args = build_parser().parse_args()
data_loader = Inferred(args.dataset[0], args.dataset[1:])
params = {
'data_loader': data_loader,
'scale': args.scale,
'dropout': args.dropout,
'backbone': args.backbone,
'learning_rate': args.learning_rate,
'optimizer': args.optimizer,
'grad_clip_norm': args.grad_clip_norm
}
num_gpus = utils.get_num_gpus()
distribution = tf.contrib.distribute.MirroredStrategy(
num_gpus=num_gpus) if num_gpus > 1 else None
config = tf.estimator.RunConfig(train_distribute=distribution,
model_dir=args.experiment,
save_summary_steps=500,
save_checkpoints_steps=500)
estimator = tf.estimator.Estimator(model_fn, params=params, config=config)
for epoch in range(args.epochs):
print('epoch {}'.format(epoch))
estimator.train(train_input_fn)
def determine_decode_and_evaluate_context(
args: argparse.Namespace, exit_stack: ExitStack,
train_context: List[mx.Context]) -> Tuple[int, Optional[mx.Context]]:
"""
Determine the number of sentences to decode and the context we should run on (CPU or GPU).
:param args: Arguments as returned by argparse.
:param exit_stack: An ExitStack from contextlib.
:param train_context: Context for training.
:return: The number of sentences to decode and a list with the context(s) to run on.
"""
num_to_decode = args.decode_and_evaluate
if args.optimized_metric == C.BLEU and num_to_decode == 0:
logger.info(
"You chose BLEU as the optimized metric, will turn on BLEU monitoring during training. "
"To control how many validation sentences are used for calculating bleu use "
"the --decode-and-evaluate argument.")
num_to_decode = -1
if num_to_decode == 0:
return 0, None
if args.use_cpu or args.decode_and_evaluate_use_cpu:
context = mx.cpu()
elif args.decode_and_evaluate_device_id is not None:
# decode device is defined from the commandline
num_gpus = utils.get_num_gpus()
check_condition(
num_gpus >= 1,
"No GPUs found, consider running on the CPU with --use-cpu "
"(note: check depends on nvidia-smi and this could also mean that the nvidia-smi "
"binary isn't on the path).")
if args.disable_device_locking:
context = utils.expand_requested_device_ids(
[args.decode_and_evaluate_device_id])
else:
context = exit_stack.enter_context(
utils.acquire_gpus([args.decode_and_evaluate_device_id],
lock_dir=args.lock_dir))
context = mx.gpu(context[0])
else:
# default decode context is the last training device
context = train_context[-1]
logger.info("Decode and Evaluate Device(s): %s", context)
return num_to_decode, context
def _setup_context(args, exit_stack):
if args.use_cpu:
context = mx.cpu()
else:
num_gpus = get_num_gpus()
check_condition(num_gpus >= 1,
"No GPUs found, consider running on the CPU with --use-cpu "
"(note: check depends on nvidia-smi and this could also mean that the nvidia-smi "
"binary isn't on the path).")
check_condition(len(args.device_ids) == 1, "cannot run on multiple devices for now")
gpu_id = args.device_ids[0]
if args.disable_device_locking:
if gpu_id < 0:
# without locking and a negative device id we just take the first device
gpu_id = 0
else:
gpu_ids = exit_stack.enter_context(acquire_gpus([gpu_id], lock_dir=args.lock_dir))
gpu_id = gpu_ids[0]
context = mx.gpu(gpu_id)
return context
def get_cuda_devices():
if "CUDA_VISIBLE_DEVICES" in os.environ:
return os.environ["CUDA_VISIBLE_DEVICES"].split(",")
else:
ngpus = get_num_gpus()
return list(range(ngpus))
def empty_dataframe():
import cudf
return cudf.DataFrame({"a": [1.0], "b": [1.0]}).head(0)
def cupy_obj():
import cupy
size = 10 ** 8
return cupy.arange(size)
@pytest.mark.slow
@pytest.mark.skipif(
get_num_gpus() <= 2, reason="Machine does not have more than two GPUs"
)
@pytest.mark.parametrize(
"cuda_obj_generator", [dataframe, empty_dataframe, series, cupy_obj]
)
@pytest.mark.parametrize("comm_api", ["tag", "am"])
def test_send_recv_cu(cuda_obj_generator, comm_api):
if comm_api == "am" and not ucp._libs.ucx_api.is_am_supported():
pytest.skip("AM only supported in UCX >= 1.11")
base_env = os.environ
env_client = base_env.copy()
# grab first two devices
cvd = get_cuda_devices()[:2]
cvd = ",".join(map(str, cvd))
# reverse CVD for other worker
def empty_dataframe():
import cudf
return cudf.DataFrame({"a": [1.0], "b": [1.0]}).head(0)
def cupy_obj():
import cupy
size = 10**8
return cupy.arange(size)
@pytest.mark.slow
@pytest.mark.skipif(get_num_gpus() <= 2,
reason="Machine does not have more than two GPUs")
@pytest.mark.parametrize("cuda_obj_generator",
[dataframe, empty_dataframe, series, cupy_obj])
@pytest.mark.parametrize("comm_api", ["tag", "am"])
def test_send_recv_cu(cuda_obj_generator, comm_api):
base_env = os.environ
env_client = base_env.copy()
# grab first two devices
cvd = get_cuda_devices()[:2]
cvd = ",".join(map(str, cvd))
# reverse CVD for other worker
env_client["CUDA_VISIBLE_DEVICES"] = cvd[::-1]
port = random.randint(13000, 15500)
# serialize function and send to the client and server
def __init__(self, worker_type, sched_addr, sched_port, worker_port,
num_gpus, run_dir, data_dir, checkpoint_dir, use_mps):
logger = logging.getLogger('worker')
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setFormatter(
logging.Formatter(LOG_FORMAT, datefmt=DATE_FORMAT, style='{'))
logger.addHandler(ch)
self._logger = logger
self._logging_handler = ch
num_available_gpus = utils.get_num_gpus()
if num_gpus > num_available_gpus:
raise ValueError('%d GPUs requested active, but only %d total '
'GPUs are available' %
(num_gpus, num_available_gpus))
signal.signal(signal.SIGINT, self._signal_handler)
self._gpu_ids = list(range(num_gpus))
self._worker_type = worker_type
self._worker_addr = socket.gethostbyname(socket.gethostname())
self._worker_port = worker_port
self._worker_rpc_client = worker_client.WorkerRpcClient(
self._worker_type, self._worker_addr, self._worker_port,
sched_addr, sched_port)
callbacks = {
'RunJob': self._run_job_callback,
'KillJob': self._kill_job_callback,
'Reset': self._reset_callback,
'Shutdown': self._shutdown_callback,
}
self._server_thread = threading.Thread(target=worker_server.serve,
args=(
worker_port,
callbacks,
))
self._server_thread.daemon = True
self._server_thread.start()
self._worker_ids, self._round_duration, error = \
self._worker_rpc_client.register_worker(len(self._gpu_ids))
if error:
raise RuntimeError(error)
if not os.path.isdir(checkpoint_dir):
# Set up a new checkpoint directory if does not already exist.
os.mkdir(checkpoint_dir)
else:
# Clear the checkpoints if they have already been created.
for dirname in os.listdir(checkpoint_dir):
if os.path.isdir(os.path.join(checkpoint_dir, dirname)):
shutil.rmtree(os.path.join(checkpoint_dir, dirname))
self._dispatcher = dispatcher.Dispatcher(self._round_duration,
self._gpu_ids,
self._worker_rpc_client,
sched_addr,
sched_port,
run_dir,
data_dir,
checkpoint_dir,
use_mps=use_mps)
self._server_thread.join()