init commit

qlib/contrib/online/manager.py

@@ -0,0 +1,147 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+import pickle
+import yaml
+import pathlib
+import pandas as pd
+import shutil
+from ..backtest.account import Account
+from ..backtest.exchange import Exchange
+from .user import User
+from .utils import load_instance
+from .utils import save_instance, init_instance_by_config
+
+
+class UserManager:
+    def __init__(self, user_data_path, save_report=True):
+        """
+        This module is designed to manager the users in online system
+        all users' data were assumed to be saved in user_data_path
+            Parameter
+                user_data_path : string
+                    data path that all users' data were saved in
+
+        variables:
+            data_path : string
+                data path that all users' data were saved in
+            users_file : string
+                A path of the file record the add_date of users
+            save_report : bool
+                whether to save report after each trading process
+            users : dict{}
+                [user_id]->User()
+                the python dict save instances of User() for each user_id
+            user_record : pd.Dataframe
+                user_id(string), add_date(string)
+                indicate the add_date for each users
+        """
+        self.data_path = pathlib.Path(user_data_path)
+        self.users_file = self.data_path / "users.csv"
+        self.save_report = save_report
+        self.users = {}
+        self.user_record = None
+
+    def load_users(self):
+        """
+        load all users' data into manager
+        """
+        self.users = {}
+        self.user_record = pd.read_csv(self.users_file, index_col=0)
+        for user_id in self.user_record.index:
+            self.users[user_id] = self.load_user(user_id)
+
+    def load_user(self, user_id):
+        """
+        return a instance of User() represents a user to be processed
+            Parameter
+                user_id : string
+            :return
+                user : User()
+        """
+        account_path = self.data_path / user_id
+        strategy_file = self.data_path / user_id / "strategy_{}.pickle".format(user_id)
+        model_file = self.data_path / user_id / "model_{}.pickle".format(user_id)
+        cur_user_list = [user_id for user_id in self.users]
+        if user_id in cur_user_list:
+            raise ValueError("User {} has been loaded".format(user_id))
+        else:
+            trade_account = Account(0)
+            trade_account.load_account(account_path)
+            strategy = load_instance(strategy_file)
+            model = load_instance(model_file)
+            user = User(account=trade_account, strategy=strategy, model=model)
+            return user
+
+    def save_user_data(self, user_id):
+        """
+        save a instance of User() to user data path
+            Parameter
+                user_id : string
+        """
+        if not user_id in self.users:
+            raise ValueError("Cannot find user {}".format(user_id))
+        self.users[user_id].account.save_account(self.data_path / user_id)
+        save_instance(
+            self.users[user_id].strategy,
+            self.data_path / user_id / "strategy_{}.pickle".format(user_id),
+        )
+        save_instance(
+            self.users[user_id].model,
+            self.data_path / user_id / "model_{}.pickle".format(user_id),
+        )
+
+    def add_user(self, user_id, config_file, add_date):
+        """
+        add the new user {user_id} into user data
+        will create a new folder named "{user_id}" in user data path
+            Parameter
+                user_id : string
+                init_cash : int
+                config_file : str/pathlib.Path()
+                   path of config file
+        """
+        config_file = pathlib.Path(config_file)
+        if not config_file.exists():
+            raise ValueError("Cannot find config file {}".format(config_file))
+        user_path = self.data_path / user_id
+        if user_path.exists():
+            raise ValueError("User data for {} already exists".format(user_id))
+
+        with config_file.open("r") as fp:
+            config = yaml.load(fp)
+        # load model
+        model = init_instance_by_config(config["model"])
+
+        # load strategy
+        strategy = init_instance_by_config(config["strategy"])
+        init_args = strategy.get_init_args_from_model(model, add_date)
+        strategy.init(**init_args)
+
+        # init Account
+        trade_account = Account(init_cash=config["init_cash"])
+
+        # save user
+        user_path.mkdir()
+        save_instance(model, self.data_path / user_id / "model_{}.pickle".format(user_id))
+        save_instance(strategy, self.data_path / user_id / "strategy_{}.pickle".format(user_id))
+        trade_account.save_account(self.data_path / user_id)
+        user_record = pd.read_csv(self.users_file, index_col=0)
+        user_record.loc[user_id] = [add_date]
+        user_record.to_csv(self.users_file)
+
+    def remove_user(self, user_id):
+        """
+        remove user {user_id} in current user dataset
+        will delete the folder "{user_id}" in user data path
+            :param
+                user_id : string
+        """
+        user_path = self.data_path / user_id
+        if not user_path.exists():
+            raise ValueError("Cannot find user data {}".format(user_id))
+        shutil.rmtree(user_path)
+        user_record = pd.read_csv(self.users_file, index_col=0)
+        user_record.drop([user_id], inplace=True)
+        user_record.to_csv(self.users_file)