import logging
import os
import time
from abc import ABCMeta
from ..config import ConfigurationManager
[docs]
class Node(metaclass=ABCMeta):
"""the basic of a node in federated learning network.
This class should be inherited instead of directly instantiate.
Attributes:
name (str): the name of this node instance.
fed_model (FedStrategyInterface): federated strategy instance
constructed according to the given configurations.
"""
[docs]
def _init_logger(self, logger_name: str, log_dir_name: str):
# TODO(fgh): move log-related into a standalone module
self.logger = logging.getLogger(logger_name)
self.logger.setLevel(logging.INFO)
time_str = time.strftime('%Y_%m%d_%H%M%S', time.localtime())
self.log_dir = os.path.join(
ConfigurationManager().log_dir_path, log_dir_name, time_str)
os.makedirs(self.log_dir, exist_ok=True)
log_file = os.path.join(self.log_dir, 'train.log')
fh = logging.FileHandler(log_file, encoding='utf8')
fh.setLevel(logging.INFO)
# create console handler with a higher log level
ch = logging.StreamHandler()
ch.setLevel(logging.ERROR)
# create formatter and add it to the handlers
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
fh.setFormatter(formatter)
ch.setFormatter(formatter)
# add the handlers to the logger
self.logger.addHandler(fh)
self.logger.addHandler(ch)
# self.fed_model.set_logger(self.logger)
[docs]
@staticmethod
def config_gpu(container_id=None):
import tensorflow as tf
cfg_mgr = ConfigurationManager()
if cfg_mgr.runtime_config.gpu_enabled:
# Please set CUDA_VISIBLE_DEVICES if not using docker
CUDA_VISIBLE_DEVICES = os.environ.get('CUDA_VISIBLE_DEVICES', '').split(',')
if len(CUDA_VISIBLE_DEVICES) > 1:
if container_id is not None:
selected_gpu = int(container_id) % len(CUDA_VISIBLE_DEVICES)
os.environ['CUDA_VISIBLE_DEVICES'] = CUDA_VISIBLE_DEVICES[selected_gpu]
else:
os.environ['CUDA_VISIBLE_DEVICES'] = CUDA_VISIBLE_DEVICES[0]
gpus = tf.config.list_physical_devices('GPU')
if gpus:
try:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
logical_devices = tf.config.list_logical_devices('GPU')
print(
len(gpus), "Physical GPUs,", len(logical_devices), "Logical GPUs"
)
except RuntimeError as e:
# Memory growth must be set before GPUs have been initialized
print(e) # TODO(fgh) expose this exception
else:
os.environ['CUDA_VISIBLE_DEVICES'] = "-1"