models.py

# -*- coding: utf-8 -*-
import os
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
import time
import queue
import datetime
import threading

import h5py
import psutil
import numpy as np

from evaluate import evaluate_ranking_metrics


DB = {'kakao_reco_730m': './tmp/kakao_reco_730m.h5py',
      'ml20m': './tmp/ml20m.h5py',
      'ml100k': './tmp/ml100k.h5py',
      'kakao_brunch_12m': './tmp/kakao_brunch_12m.h5py'}


def collect_memory_usage(stop_event, result_queue):
    p = psutil.Process(os.getpid())
    data = []
    while not stop_event.is_set():
        data.append(p.memory_info().rss)
        time.sleep(5)
    result_queue.put(data)


def db_to_coo(db):
    from scipy.sparse import coo_matrix
    V = db.attrs['num_items']
    U = db.attrs['num_users']
    col = db['rowwise']['key'][::]
    data = db['rowwise']['val'][::]
    indptr = db['rowwise']['indptr'][::]
    real_size = indptr[-1]
    col = col[:real_size]
    data = data[:real_size]
    row = np.zeros(col.shape, dtype=np.int32)
    start, row_id = 0, 0
    for end in indptr:
        row[start:end] = row_id
        row_id += 1
        start = end
    coo = coo_matrix((data, (row, col)), (U, V))
    data, row, col = None, None, None
    return coo


def db_to_dataframe(db, spark, context):
    from pyspark.sql import Row
    coo = db_to_coo(db)
    data = context.parallelize(np.array([coo.row, coo.col, coo.data]).T)
    coo = None
    return spark.createDataFrame(data.map(lambda p: Row(row=int(p[0]),
                                                        col=int(p[1]),
                                                        data=float(p[2]))))


def get_buffalo_db(db):
    from buffalo.data.mm import MatrixMarket
    db_opt = BuffaloLib().get_database(db)
    db = MatrixMarket(db_opt)
    db.create()
    return db


class Benchmark(object):
    def get_option(self, lib_name, algo_name, **kwargs):
        if lib_name == 'buffalo':
            if algo_name == 'als':
                from buffalo.algo.options import ALSOption
                opt = ALSOption().get_default_option()
                opt.update({'d': kwargs.get('d', 100),
                            'optimizer': {True: 'manual_cg', False: 'ldlt'}.get(kwargs.get('use_cg', True)),
                            'num_iters': kwargs.get('num_iters', 10),
                            'num_cg_max_iters': 3,
                            'validation': kwargs.get('validation'),
                            'accelerator': kwargs.get('gpu', False),
                            'num_workers': kwargs.get('num_workers', 10),
                            'compute_loss_on_training': kwargs.get('compute_loss_on_training', False)})
                return opt
            if algo_name == 'bpr':
                from buffalo.algo.options import BPRMFOption
                opt = BPRMFOption().get_default_option()
                opt.update({'d': kwargs.get('d', 100),
                            'validation': kwargs.get('validation'),
                            'num_iters': kwargs.get('num_iters', 10),
                            'num_workers': kwargs.get('num_workers', 10),
                            'compute_loss_on_training': kwargs.get('compute_loss_on_training', False)})
                return opt
            if algo_name == 'warp':
                from buffalo.algo.options import WARPOption
                opt = WARPOption().get_default_option()
                opt.update({'d': kwargs.get('d', 100),
                            'validation': kwargs.get('validation'),
                            'num_iters': kwargs.get('num_iters', 10),
                            'max_trials': 100,
                            'num_workers': kwargs.get('num_workers', 10),
                            'compute_loss_on_training': kwargs.get('compute_loss_on_training', False)})
                return opt
        elif lib_name == 'implicit':
            if algo_name == 'als':
                return {'factors': kwargs.get('d', 100),
                        'dtype': np.float32,
                        'use_native': True,
                        'use_gpu': kwargs.get('gpu', False),
                        'use_cg': kwargs.get('use_cg', True),
                        'iterations': kwargs.get('num_iters', 10),
                        'num_threads': kwargs.get('num_workers', 10),
                        'calculate_training_loss': kwargs.get('calculate_training_loss', False)}
            if algo_name == 'bpr':
                return {'factors': kwargs.get('d', 100),
                        'dtype': np.float32,
                        'iterations': kwargs.get('num_iters', 10),
                        'verify_negative_samples': True,
                        'num_threads': kwargs.get('num_workers', 10)}
        elif lib_name == 'lightfm':
            if algo_name == 'bpr':
                return {'epochs': kwargs.get('num_iters', 10),
                        'verbose': True,
                        'num_threads': kwargs.get('num_workers', 10)}
            if algo_name == 'warp':
                return {'epochs': kwargs.get('num_iters', 10),
                        'verbose': True,
                        'num_threads': kwargs.get('num_workers', 10)}
        elif lib_name == 'pyspark':
            if algo_name == 'als':
                return {'maxIter': kwargs.get('num_iters', 10),
                        'rank': kwargs.get('d', 100),
                        'alpha': 8,
                        'implicitPrefs': True,
                        'userCol': 'row',
                        'itemCol': 'col',
                        'intermediateStorageLevel': 'MEMORY_ONLY',
                        'finalStorageLevel': 'MEMORY_ONLY',
                        'ratingCol': 'data'}

    def run(self, func, *args, **kwargs):
        stop_event = threading.Event()
        result_queue = queue.Queue()
        t = threading.Thread(target=collect_memory_usage, args=(stop_event, result_queue))
        t.start()
        time.sleep(0.1)  # context switching
        start_t = time.time()
        if kwargs.get('iterable'):
            iterable = kwargs.get('iterable')
            kwargs.pop('iterable')
            for data in iterable:
                func(data, **kwargs)
        else:
            func(*args, **kwargs)
        elapsed = time.time() - start_t
        stop_event.set()
        t.join()
        memory_usage = result_queue.get(block=True, timeout=10)
        return elapsed, {'min': min(memory_usage) / 1024 / 1024.0,
                         'avg': sum(memory_usage) / len(memory_usage) / 1024 / 1024.0,
                         'max': max(memory_usage) / 1024 / 1024.0}


class ImplicitLib(Benchmark):
    def __init__(self):
        super().__init__()

    def get_database(self, name, **kwargs):
        if name in ['ml20m', 'ml100k', 'kakao_reco_730m', 'kakao_brunch_12m']:
            db = h5py.File(DB[name])
            ratings = db_to_coo(db)
            db.close()
            return ratings

    def als(self, database, **kwargs):
        from implicit.als import AlternatingLeastSquares
        opts = self.get_option('implicit', 'als', **kwargs)
        model = AlternatingLeastSquares(
            **opts
        )
        ratings = self.get_database(database, **kwargs)
        if kwargs.get('return_instance_before_train'):
            return (model, ratings)

        elapsed, mem_info = self.run(model.fit, ratings)
        if kwargs.get('return_instance'):
            return model
        model = None
        return elapsed, mem_info

    def bpr(self, database, **kwargs):
        from implicit.bpr import BayesianPersonalizedRanking
        opts = self.get_option('implicit', 'bpr', **kwargs)
        model = BayesianPersonalizedRanking(
            **opts
        )
        ratings = self.get_database(database, **kwargs)
        if kwargs.get('return_instance_before_train'):
            return (model, ratings)

        elapsed, mem_info = self.run(model.fit, ratings)
        if kwargs.get('return_instance'):
            return model
        model = None
        return elapsed, mem_info

    def most_similar(self, keys, **kwargs):
        model = kwargs['model']
        kwargs.pop('model')
        elapsed, mem_info = self.run(model.similar_items, keys, **kwargs)
        return elapsed, mem_info


class BuffaloLib(Benchmark):
    def __init__(self):
        super().__init__()

    def get_database(self, name, **kwargs):
        from buffalo.data.mm import MatrixMarketOptions
        data_opt = MatrixMarketOptions().get_default_option()
        if kwargs.get('validation', None) is None:
            data_opt.validation = None
        data_opt.data.use_cache = True
        data_opt.data.batch_mb = kwargs.get('batch_mb', 1024)
        if name == 'ml20m':
            data_opt.data.path = DB[name]
            data_opt.input.main = '../tests/ext/ml-20m/main'
        elif name =='ml100k':
            data_opt.data.path = DB[name]
            data_opt.input.main = '../tests/ext/ml-100k/main'
        elif name == 'kakao_reco_730m':
            data_opt.data.path = DB[name]
            data_opt.data.tmp_dir = './tmp/'
            data_opt.input.main = '../tests/ext/kakao-reco-730m/main'
        elif name == 'kakao_brunch_12m':
            data_opt.data.path = DB[name]
            data_opt.data.tmp_dir = './tmp/'
            data_opt.input.main = '../tests/ext/kakao-brunch-12m/main'
        return data_opt

    def als(self, database, **kwargs):
        from buffalo.algo.als import ALS
        opts = self.get_option('buffalo', 'als', **kwargs)
        data_opt = self.get_database(database, **kwargs)
        als = ALS(opts, data_opt=data_opt)
        als.initialize()
        if kwargs.get('return_instance_before_train'):
            return als
        elapsed, mem_info = self.run(als.train)
        if kwargs.get('return_instance'):
            return als
        als = None
        return elapsed, mem_info

    def validation(self, algo_name, database, **kwargs):
        inst = getattr(self, algo_name)(database, return_instance=True, **kwargs)
        ret = inst.get_validation_results()  # same as below
        for p in ['error', 'rmse']:
            ret.pop(p)
        return ret
        inst.data._prepare_validation_data()
        K = kwargs.get('validation', {}).get('topk', 10)
        userids = list(set(inst.data.handle['vali']['row'][::]))
        itemids = list(range(inst.data.handle['idmap']['cols'].shape[0]))
        recs = []
        for user in userids:
            seen = inst.data.vali_data['validation_seen'].get(user, set())
            user_str = inst.data.handle['idmap']['rows'][user].decode('utf-8')
            topn = inst.topk_recommendation(user_str, topk=K + len(seen))
            topn = [inst._idmanager.itemid_map[t] for t in topn]
            topn = [t for t in topn if t not in seen][:K]
            recs.append((user, topn))
        return evaluate_ranking_metrics(recs, K, inst.data.vali_data, inst.data.header['num_items'])

    def bpr(self, database, **kwargs):
        from buffalo.algo.bpr import BPRMF
        opts = self.get_option('buffalo', 'bpr', **kwargs)
        data_opt = self.get_database(database, **kwargs)
        bpr = BPRMF(opts, data_opt=data_opt)
        bpr.initialize()
        if kwargs.get('return_instance_before_train'):
            return bpr
        elapsed, mem_info = self.run(bpr.train)
        if kwargs.get('return_instance'):
            return bpr
        bpr = None
        return elapsed, mem_info

    def warp(self, database, **kwargs):
        from buffalo.algo.warp import WARP
        opts = self.get_option('buffalo', 'warp', **kwargs)
        data_opt = self.get_database(database, **kwargs)
        warp = WARP(opts, data_opt=data_opt)
        warp.initialize()
        if kwargs.get('return_instance_before_train'):
            return warp
        elapsed, mem_info = self.run(warp.train)
        if kwargs.get('return_instance'):
            return warp
        warp = None
        return elapsed, mem_info

    def most_similar(self, keys, **kwargs):
        model = kwargs['model']
        kwargs.pop('model')
        elapsed, mem_info = self.run(model.most_similar, keys, **kwargs)
        return elapsed, mem_info


class LightfmLib(Benchmark):
    def __init__(self):
        super().__init__()

    def get_database(self, name, **kwargs):
        if name in ['ml20m', 'ml100k', 'kakao_reco_730m', 'kakao_brunch_12m']:
            db = h5py.File(DB[name])
            ratings = db_to_coo(db)
            db.close()
            return ratings

    def als(self, database, **kwargs):
        raise NotImplementedError

    def bpr(self, database, **kwargs):
        from lightfm import LightFM
        opts = self.get_option('lightfm', 'bpr', **kwargs)
        data = self.get_database(database, **kwargs)
        bpr = LightFM(loss='bpr',
                      no_components=kwargs.get('num_workers'))
        elapsed, mem_info = self.run(bpr.fit, data, data, **opts)
        if kwargs.get('return_instance'):
            return bpr
        bpr = None
        return elapsed, mem_info

    def validation(self, algo_name, database, **kwargs):
        K = kwargs.get('validation', {}).get('topk', 10)
        inst = getattr(self, algo_name)(database, return_instance=True, **kwargs)
        db = get_buffalo_db(database)
        db._prepare_validation_data()
        userids = list(set(db.handle['vali']['row'][::]))
        itemids = list(range(db.handle['idmap']['cols'].shape[0]))
        recs = []
        for user in userids:
            topn = np.argsort(-inst.predict(user, itemids))
            topn = [t for t in topn if t not in db.vali_data['validation_seen'].get(user, set())][:K]
            recs.append((user, topn))
        return evaluate_ranking_metrics(recs, K, db.vali_data, db.header['num_items'])

    def warp(self, database, **kwargs):
        from lightfm import LightFM
        opts = self.get_option('lightfm', 'warp', **kwargs)
        data = self.get_database(database, **kwargs)
        warp = LightFM(loss='warp',
                       learning_schedule='adagrad',
                       no_components=kwargs.get('num_workers'),
                       max_sampled=100)
        elapsed, mem_info = self.run(warp.fit, data, data, **opts)
        if kwargs.get('return_instance'):
            return warp
        bpr = None
        return elapsed, mem_info


class QmfLib(Benchmark):
    def __init__(self):
        super().__init__()
        self.bin_root = '../../qmf.git/obj/bin'

    def get_database(self, name, **kwargs):
        if name in ['ml20m', 'ml100k']:
            db_path = f'./qmf.{name}.dataset'
            num_header_lines = 4
            if not os.path.isfile(db_path):
                with open('../tests/%s/main' % {'ml20m': 'ml-20m', 'ml100k': 'ml-100k'}.get(name)) as fin:
                    for i in range(num_header_lines):
                        _ = fin.readline()
                    with open(db_path, 'w') as fout:
                        for line in fin:
                            fout.write(line)
            return os.path.abspath(db_path)
        if name in ['kakao_brunch_12m']:
            db_path = f'./qmf.{name}.dataset'
            num_header_lines = 4
            if not os.path.isfile(db_path):
                with open('../tests/ext/kakao-brunch-12m/main') as fin:
                    for i in range(num_header_lines):
                        _ = fin.readline()
                    with open(db_path, 'w') as fout:
                        for line in fin:
                            fout.write(line)
            return os.path.abspath(db_path)
        elif name == 'kakao_reco_730m':
            db_path = f'./qmf.{name}.dataset'
            num_header_lines = 4
            if not os.path.isfile(db_path):
                with open('../tests/ext/kakao-reco-730m/main') as fin:
                    for i in range(num_header_lines):
                        _ = fin.readline()
                    with open(db_path, 'w') as fout:
                        for line in fin:
                            fout.write(line)
            return os.path.abspath(db_path)

    def als(self, database, **kwargs):
        import subprocess
        train_ds = self.get_database(database, **kwargs)
        d = kwargs.get('d')
        num_iters = kwargs.get('num_iters', 10)
        num_workers = kwargs.get('num_workers')
        cmd = ['./wals',
               f'--train_dataset={train_ds}',
               '--user_factors=/dev/null',
               '--item_factors=/dev/null',
               f'--nfactors={d}',
               f'--nepochs={num_iters}',
               f'--nthreads={num_workers}']
        stop_event = threading.Event()
        result_queue = queue.Queue()
        t = threading.Thread(target=collect_memory_usage, args=(stop_event, result_queue))
        t.start()
        ret = subprocess.run(cmd, cwd=self.bin_root, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stop_event.set()
        t.join()
        memory_usage = result_queue.get()

        start_t, end_t = None, None
        for line in ret.stderr.decode('utf-8').split('\n'):
            if line.endswith('training'):
                hms = line.split()[1]
                start_t = datetime.datetime.strptime(hms, '%H:%M:%S.%f')
            elif line.endswith('saving model output'):
                hms = line.split()[1]
                end_t = datetime.datetime.strptime(hms, '%H:%M:%S.%f')
        elapsed = (end_t - start_t).seconds + (end_t - start_t).microseconds / 1000 / 1000
        return elapsed, {'min': min(memory_usage) / 1024 / 1024.0,
                         'avg': sum(memory_usage) / len(memory_usage) / 1024 / 1024.0,
                         'max': max(memory_usage) / 1024 / 1024.0}

    def bpr(self, database, **kwargs):
        import subprocess
        train_ds = self.get_database(database, **kwargs)
        d = kwargs.get('d')
        num_iters = kwargs.get('num_iters', 10)
        num_workers = kwargs.get('num_workers')
        cmd = ['./bpr',
               f'--train_dataset={train_ds}',
               '--user_factors=/dev/null',
               '--item_factors=/dev/null',
               '--num_negative_samples=1',
               '--eval_num_neg=0',
               f'--nfactors={d}',
               f'--nepochs={num_iters}',
               f'--nthreads={num_workers}']
        stop_event = threading.Event()
        result_queue = queue.Queue()
        t = threading.Thread(target=collect_memory_usage, args=(stop_event, result_queue))
        t.start()
        ret = subprocess.run(cmd, cwd=self.bin_root, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stop_event.set()
        t.join()
        memory_usage = result_queue.get()

        start_t, end_t = None, None
        for line in ret.stderr.decode('utf-8').split('\n'):
            if line.endswith('training'):
                hms = line.split()[1]
                start_t = datetime.datetime.strptime(hms, '%H:%M:%S.%f')
            elif line.endswith('saving model output'):
                hms = line.split()[1]
                end_t = datetime.datetime.strptime(hms, '%H:%M:%S.%f')
        elapsed = (end_t - start_t).seconds + (end_t - start_t).microseconds / 1000 / 1000
        return elapsed, {'min': min(memory_usage) / 1024 / 1024.0,
                         'avg': sum(memory_usage) / len(memory_usage) / 1024 / 1024.0,
                         'max': max(memory_usage) / 1024 / 1024.0}


class PysparkLib(Benchmark):
    def __init__(self):
        super().__init__()

    def get_database(self, name, **kwargs):
        if name in ['ml20m', 'ml100k', 'kakao_reco_730m', 'kakao_brunch_12m']:
            db = h5py.File(DB[name])
            ratings = db_to_dataframe(db, kwargs.get('spark'), kwargs.get('context'))
            db.close()
            return ratings

    def als(self, database, **kwargs):
        from pyspark.sql import SparkSession
        from pyspark.ml.recommendation import ALS
        from pyspark import SparkConf, SparkContext
        opts = self.get_option('pyspark', 'als', **kwargs)
        conf = SparkConf()\
               .setAppName("pyspark")\
               .setMaster('local[%s]' % kwargs.get('num_workers'))\
               .set('spark.local.dir', './tmp/')\
               .set('spark.worker.cleanup.enabled', 'true')\
               .set('spark.driver.memory', '32G')
        context = SparkContext(conf=conf)
        context.setLogLevel('WARN')
        spark = SparkSession(context)
        data = self.get_database(database, spark=spark, context=context)
        print(opts)
        als = ALS(**opts)
        elapsed, memory_usage = self.run(als.fit, data)
        spark.stop()
        return elapsed, memory_usage

    def bpr(self, database, **kwargs):
        raise NotImplementedError