<filename>built-in/TensorFlow/Official/cv/image_classification/ResnetVariant_for_TensorFlow/automl/vega/core/scheduler/local_master.py
# -*- coding: utf-8 -*-

# Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
# This program is free software; you can redistribute it and/or modify
# it under the terms of the MIT License.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# MIT License for more details.

"""The LocalMaster's method is same as Master, and the class is used on single node."""
import os
from ..trainer.utils import WorkerTypes
from vega.core.common.general import General


class LocalMaster(object):
    """The Master's method is same as Master."""

    def __init__(self):
        """Init master."""
        self.cfg = General
        self.step_name = None
        self.worker_id = None
        if os.environ['DEVICE_CATEGORY'] == 'NPU':
            os.environ['RANK_SIZE'] = '1'
            os.environ.pop('RANK_TABLE_FILE', None)

    def run(self, worker):
        """Run a worker, call the worker's train_prcess() method.

        :param worker: a worker.
        :type worker: object that the class was inherited from DistributedWorker.

        """
        if worker is None:
            return
        self.step_name = worker.step_name
        self.worker_id = worker.worker_id
        if worker.worker_type == WorkerTypes.EVALUATOR:
            for sub_worker in worker.sub_worker_list:
                sub_worker.train_process()
        elif worker.worker_type == WorkerTypes.HAVA_D_EVALUATOR:
            pass
        else:
            worker.train_process()

    def join(self):
        """Return immediately."""
        return

    def pop_finished_worker(self, train_worker=True):
        """Pop saved worker id and step name.

        :return: the finished worker info, include step_name and worker_id.
            eg. {"step_name":"round1", "worker_id":1}
        :rtype: dict or None

        """
        if self.worker_id is None:
            return None
        worker_id = self.worker_id
        step_name = self.step_name
        self.worker_id = None
        self.step_name = None
        return {"step_name": step_name, "worker_id": worker_id}

    def pop_finished_train_worker(self):
        """Pop saved worker id and step name.

        :return: the finished worker info, include step_name and worker_id.
            eg. {"step_name":"round1", "worker_id":1}
        :rtype: dict or None

        """
        return self.pop_finished_worker(train_worker=True)

    def pop_finished_evaluate_worker(self):
        """Pop saved worker id and step_name.

        :return: the finished worker info, include step_name and worker_id.
            eg. {"step_name":"round1", "worker_id":1}
        :rtype: dict or None

        """
        return self.pop_finished_worker(train_worker=False)

    def pop_all_finished_train_worker(self):
        """Pop saved worker id and step name.

        :return: a finished worker info list.
        :rtype: list of dict

        """
        worker_info_list = []
        finished_train_worker_info = self.pop_finished_train_worker()
        if finished_train_worker_info is not None:
            worker_info_list.append(finished_train_worker_info)
        return worker_info_list

    def pop_all_finished_evaluate_worker(self):
        """Pop saved worker id and step name.

        :return: a finished worker info list.
        :rtype: list of dict

        """
        return self.pop_all_finished_train_worker()

    def close_cluster(self):
        """Close all distributed cluster, implement the interface without actually closing."""
        pass

    def close_client(self):
        """Close cluster client, implement the interface without actually closing."""
        pass

    @staticmethod
    def shutdown():
        """Shut down the cluster, implement the interface without actually shutting down."""
        pass
