logginghelper

logging_helper.py

@@ -0,0 +1,205 @@
+"""
+a module of logs saving and backuping
+"""
+
+import os
+import datasets as ds
+from apscheduler.schedulers.background import BackgroundScheduler
+from tqdm import tqdm
+from utils import beijing, md5, json_to_str
+from huggingface_hub import HfApi
+import pandas as pd
+import glob
+
+hf = HfApi()
+hf.token = os.environ.get("hf_token")
+
+
+class LoggingHelper:
+
+    def __init__(
+        self,
+        repo_id: str,
+        local_dir: str = "data/logs",
+        synchronize_interval: int = 60,
+    ):
+        """
+        :param repo_id: the repo_id of the dataset in huggingface
+        :param local_dir: a directory in cwd to store downloaded files
+        :param synchronize_interval: the interval of synchronizing between local and huggingface
+
+        """
+        self.local_dir = local_dir
+        self.repo_id = repo_id
+        self.synchronize_interval = synchronize_interval
+        self.repo_type = "dataset"
+        self.scheduler = BackgroundScheduler()
+        self.buffer = dict[str, ds.Dataset]()
+        self.need_push = dict[str, bool]()
+        self.today = beijing().date()
+        ds.disable_progress_bar()
+        self.dataframe: pd.DataFrame
+        self.pull()
+        self.start_synchronize()
+
+    def addlog(self, log: dict):
+        """add one log"""
+        remotedir = self.remotedir()
+        filename = md5(json_to_str(log))[:2] + ".json"
+        remotepath = "/".join([remotedir, filename])
+        if remotepath in self.buffer:
+            self.buffer[remotepath] = self.buffer[remotepath].add_item(log)  # type: ignore
+        else:
+            self.buffer[remotepath] = ds.Dataset.from_dict({})
+            self.buffer[remotepath] = self.buffer[remotepath].add_item(log)  # type: ignore
+        self.need_push[remotepath] = True
+        print("[addlog] Added a log to buffer")
+
+    def remotedir(self):
+        now = beijing()
+        year = now.year.__str__()
+        month = now.month.__str__()
+        day = now.day.__str__()
+        return "/".join([year, month, day])
+
+    def pull(self):
+        try:
+            self.download()
+            remotedir = self.remotedir()
+            print(f"[pull] today dir: {remotedir}")
+            filenames = hf.list_repo_files(
+                repo_id=self.repo_id,
+                repo_type=self.repo_type,
+            )
+            files_to_load = [
+                filename
+                for filename in filenames
+                if filename not in self.buffer
+                and filename.startswith(remotedir)
+                and filename.endswith(".json")
+            ]
+            print(f"[pull] total {len(files_to_load)} to load")
+            for filename in tqdm(files_to_load):
+                print()
+                path = os.sep.join([self.local_dir, filename])
+                with open(path, "r") as f:
+                    data = f.read()
+                if len(data) != 0:
+                    self.buffer[filename] = ds.Dataset.from_json(path)  # type: ignore
+                    self.need_push[filename] = False
+            return True
+        except Exception as e:
+            print(f"[pull] {type(e)}: {e}")
+            return False
+
+    def push_yesterday(self) -> bool:
+        try:
+            year = self.today.year.__str__()
+            month = self.today.month.__str__()
+            day = self.today.day.__str__()
+            remotedir = "/".join([year, month, day])
+            files_to_push = []
+            for filename in self.buffer:
+                if not filename.startswith(remotedir):
+                    continue
+                if not self.need_push[filename]:
+                    del self.buffer[filename]
+                    del self.need_push[filename]
+                files_to_push.append(filename)
+            if len(files_to_push) == 0:
+                return True
+            print("[push_yesterday] Writing datasets to json files")
+            for filename in files_to_push:
+                localpath = os.sep.join([self.local_dir, filename])
+                self.buffer[filename].to_json(localpath)
+                files_to_push.append(filename)
+                print(f"[push_yesterday] {filename} finished")
+            print("[push_yesterday] Done")
+            print("[push_yesterday] Pushing log files to remote")
+            if len(files_to_push):
+                localdir = os.sep.join([self.local_dir, remotedir])
+                res = hf.upload_folder(
+                    repo_id=self.repo_id,
+                    folder_path=localdir,
+                    path_in_repo=remotedir,
+                    repo_type=self.repo_type,
+                    commit_message=f"Updated at {beijing()}",
+                )
+                print(f"[push_yesterday] Log files pushed to {res}")
+            print("[push_yesterday] Done")
+            return True
+        except Exception as e:
+            print(f"[push_yesterday] {type(e)}: {e}")
+            return False
+
+    def push(self):
+        try:
+            now = beijing().date()
+            if now != self.today:  # new day comes
+                if not self.push_yesterday():
+                    print("[push] Failed to upload yesterday's log files")
+                self.today = now
+            files_to_push = [
+                filename for filename in self.need_push if self.need_push[filename]
+            ]
+            if len(files_to_push) == 0:
+                return True
+            print("[push] Writing datasets to json files")
+            for filename in files_to_push:
+                localpath = os.path.join(self.local_dir, filename)
+                self.buffer[filename].to_json(localpath)
+                print(f"[push] {filename} finished")
+            print("[push] Done")
+            print("[push] Pushing log files to remote")
+            remotedir = self.remotedir()
+            localdir = "/".join([self.local_dir, remotedir])
+            res = hf.upload_folder(
+                repo_id=self.repo_id,
+                folder_path=localdir,
+                path_in_repo=remotedir,
+                repo_type=self.repo_type,
+                commit_message=f"Updated at {beijing()}",
+            )
+            print(f"[push] Log files pushed to {res}")
+            print("[push] Done")
+            return True
+        except Exception as e:
+            print(f"[push] {type(e)}: {e}")
+            return False
+
+    def download(self):
+        print("[download] Starting downloading")
+        try:
+            res = hf.snapshot_download(
+                repo_id=self.repo_id,
+                repo_type="dataset",
+                local_dir=self.local_dir,
+            )
+            print(f"[download] Downloaded to {res}")
+        except Exception as e:
+            print(f"[download] {type(e)}: {e}")
+        print("[download] Done")
+
+    def start_synchronize(self):
+        self.scheduler.add_job(
+            self.push,
+            "interval",
+            seconds=self.synchronize_interval,
+        )
+        self.scheduler.start()
+
+    def refresh(self) -> list[dict]:
+        self.download()
+        files = glob.glob("**/*.json", root_dir=self.local_dir, recursive=True)
+        filepathes = [os.sep.join([self.local_dir, file]) for file in files]
+        datasets = []
+        for path in tqdm(filepathes):
+            path = str(path)
+            datasets.append(ds.Dataset.from_json(path))
+        df = pd.DataFrame()
+        if datasets:
+            dataset: ds.Dataset = ds.concatenate_datasets(datasets)
+            df = dataset.to_pandas()
+            assert isinstance(df, pd.DataFrame)
+            df = df.sort_values(by="timestamp")
+        return df.to_dict(orient="records")