# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
from __future__ import print_function, absolute_import

import argparse
import os

import chainermn
import sagemaker_containers

if __name__ == '__main__':
    env = sagemaker_containers.training_env()

    parser = argparse.ArgumentParser()

    parser.add_argument('--num-gpus', type=int, default=env.num_gpus)
    parser.add_argument('--communicator', type=str,
                        default='naive' if env.num_gpus == 0 else 'pure_nccl')
    parser.add_argument('--hosts', type=str, default=env.hosts)
    parser.add_argument('--node_to_fail', type=int)
    parser.add_argument('--output-data-dir', type=str, default=env.output_data_dir)

    args = parser.parse_args()

    if len(args.hosts) == 1:
        open(os.path.join(env.output_data_dir, 'this_file_is_expected'), 'a').close()
        raise Exception('intentional exception on a single machine')

    comm = chainermn.create_communicator(args.communicator)

    # When running in local mode, setting rank to 'inter_rank' simulates multi-node training.
    rank = comm.inter_rank

    if args.node_to_fail == rank:
        print('intentional exception from node {}'.format(rank))
        open(os.path.join(env.output_data_dir, 'file_from_node_{}'.format(rank)), 'a').close()
        raise Exception('intentional exception from node {}'.format(rank))