### This file is originally from: [mlcommons repo](https://github.com/mlcommons/training/tree/9947bdf21ee3f2488fa4b362eec2ce7deb2ec4dd/single_stage_detector/ssd/mlperf_logger.py)
# Copyright (c) 2018, NVIDIA CORPORATION. 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.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

import torch
import numpy as np
import os
from mlperf_logging import mllog
from mlperf_logging.mllog import constants as mllog_const
mllogger = mllog.get_mllogger()
mllog.config(
    filename=(os.getenv("COMPLIANCE_FILE") or "mlperf_compliance.log"),
    root_dir=os.path.normpath(os.path.dirname(os.path.realpath(__file__))))

def ssd_print(*args, sync=True, **kwargs):
    use_cuda = os.getenv('USE_CUDA')
    if sync and use_cuda=='True':
        barrier()
    if get_rank() == 0:
        kwargs['stack_offset'] = 2
        mllogger.event(*args, **kwargs)


def barrier():
    """
    Works as a temporary distributed barrier, currently pytorch
    doesn't implement barrier for NCCL backend.
    Calls all_reduce on dummy tensor and synchronizes with GPU.
    """
    if torch.distributed.is_initialized():
        torch.distributed.all_reduce(torch.cuda.FloatTensor(1))
        torch.cuda.synchronize()


def get_rank():
    """
    Gets distributed rank or returns zero if distributed is not initialized.
    """
    if torch.distributed.is_initialized():
        rank = torch.distributed.get_rank()
    else:
        rank = os.getenv('RANK', os.getenv('LOCAL_RANK', 0))
    return rank

def broadcast_seeds(seed, device):
    if torch.distributed.is_initialized():
        seeds_tensor = torch.LongTensor([seed]).to(device)
        torch.distributed.broadcast(seeds_tensor, 0)
        seed = seeds_tensor.item()
    return seed

[((21, 11, 21, 31), 'mlperf_logging.mllog.get_mllogger', 'mllog.get_mllogger', ({}, {}), '()', False, 'from mlperf_logging import mllog\n'), ((27, 15, 27, 36), 'os.getenv', 'os.getenv', ({(27, 25, 27, 35): '"""USE_CUDA"""'}, {}), "('USE_CUDA')", False, 'import os\n'), ((41, 7, 41, 41), 'torch.distributed.is_initialized', 'torch.distributed.is_initialized', ({}, {}), '()', False, 'import torch\n'), ((50, 7, 50, 41), 'torch.distributed.is_initialized', 'torch.distributed.is_initialized', ({}, {}), '()', False, 'import torch\n'), ((57, 7, 57, 41), 'torch.distributed.is_initialized', 'torch.distributed.is_initialized', ({}, {}), '()', False, 'import torch\n'), ((43, 8, 43, 32), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n'), ((51, 15, 51, 43), 'torch.distributed.get_rank', 'torch.distributed.get_rank', ({}, {}), '()', False, 'import torch\n'), ((59, 8, 59, 52), 'torch.distributed.broadcast', 'torch.distributed.broadcast', ({(59, 36, 59, 48): 'seeds_tensor', (59, 50, 59, 51): '(0)'}, {}), '(seeds_tensor, 0)', False, 'import torch\n'), ((23, 14, 23, 42), 'os.getenv', 'os.getenv', ({(23, 24, 23, 41): '"""COMPLIANCE_FILE"""'}, {}), "('COMPLIANCE_FILE')", False, 'import os\n'), ((42, 37, 42, 62), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(42, 60, 42, 61): '(1)'}, {}), '(1)', False, 'import torch\n'), ((53, 33, 53, 59), 'os.getenv', 'os.getenv', ({(53, 43, 53, 55): '"""LOCAL_RANK"""', (53, 57, 53, 58): '0'}, {}), "('LOCAL_RANK', 0)", False, 'import os\n'), ((24, 46, 24, 72), 'os.path.realpath', 'os.path.realpath', ({(24, 63, 24, 71): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((58, 23, 58, 47), 'torch.LongTensor', 'torch.LongTensor', ({(58, 40, 58, 46): '[seed]'}, {}), '([seed])', False, 'import torch\n')]