import logging  # 引入logging模块
import os.path
import time
# 第一步，创建一个logger
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # Log等级总开关
# 第二步，创建一个handler，用于写入日志文件
rq = time.strftime('%Y%m%d%H%M', time.localtime(time.time()))
log_path = os.path.dirname(os.getcwd()) + '/Logs/'
log_name = log_path + rq + '.log'
logfile = log_name
fh = logging.FileHandler(logfile, mode='w')
fh.setLevel(logging.DEBUG)  # 输出到file的log等级的开关
# 第三步，定义handler的输出格式
formatter = logging.Formatter("%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s")
fh.setFormatter(formatter)
# 第四步，将logger添加到handler里面
logger.addHandler(fh)

import argparse
import redis
from tqdm import tqdm
from multiprocessing import Pool, Process

parser = argparse.ArgumentParser(description='PyTorch MM Training')
parser.add_argument('--port', default=6379, type=str, help="port id")
parser.add_argument('--file', default=None, type=str, help="paths for train instances")
args = parser.parse_args()

host = '127.0.0.1'
port = args.port

# r = redis.Redis(host=host, port=port)

train_file = args.file

def worker_i(train_file, i, nums=16):
    r = redis.Redis(host=host, port=port)
    k = 0

    if i == 0:
        iter_ = tqdm(open(train_file, 'r'))
    else:
        iter_ = open(train_file, 'r')

    for info in iter_:
        if k % nums == i: # 将多个任务分成nums份
            r.set(k, info.strip())
        k += 1
process_list = []

nums = 8
for i in range(nums):
    p = Process(target=worker_i, args=(train_file, i, nums))
    p.start()
    process_list.append(p)

for i in process_list:
    p.join()

print('写入完成')