src/testing.py

from __future__ import annotations  # For self-referencing annotations

import json
import os
import shutil
import sqlite3
import sys

from huggingface_hub import Repository
from random import choices
from typing import List, Dict, Optional

from src.architectures import Architecture
from src.common import data_dir


class TestGenerator:
    """
    Wrapper class to hold testing questions and serve up examples
    """
    questions: List[str] = None

    @classmethod
    def load_questions(cls, reload=False) -> None:
        """
        Load the available questions from the json file.
        Default to not re-loading if already done, but allow for the option to do so
        """
        if cls.questions is not None and not reload:
            return
        question_file = os.path.join(data_dir, 'json', 'test_questions.json')
        with open(question_file, 'r') as f:
            question_json = json.load(f)
        cls.questions = question_json['questions']

    @classmethod
    def question_count(cls) -> int:
        cls.load_questions()
        return len(cls.questions)

    @classmethod
    def get_random_questions(cls, n: int):
        """
        Return n random questions
        """
        cls.load_questions()
        return choices(cls.questions, k=n)


class ArchitectureRequestRecord:
    """
    Representation of the test data associated with each invocation of an architecture
    """
    all: List[ArchitectureRequestRecord] = None

    class ArchStep:
        """
        Inner class to just hold this data
        """
        def __init__(self, name: str, start: int, end: int):
            self.name = name
            self.start = start
            self.end = end
            self.elapsed = end - start

    def __init__(self, arch: str, response_len: int, start: int, end: int,
                 elapsed: int, tags: List[str], test_group: Optional[str],
                 comment: str, steps: List[ArchitectureRequestRecord.ArchStep]):
        self.arch = arch
        self.response_len = response_len
        self.start = start
        self.end = end
        self.elapsed = elapsed
        self.tags = tags
        self.test_group = test_group
        self.comment = comment
        self.steps = steps

    @classmethod
    def from_dict(cls, test: Dict) -> ArchitectureRequestRecord:
        arch = test['architecture']
        response_len = len(test['request']['response_evolution'][-1])
        start = test['trace']['steps'][0]['start_ms']
        end = test['trace']['steps'][-1]['end_ms']
        elapsed = end - start
        tags = test['test_tags']
        test_group = None
        for tag in tags:
            if tag.startswith("TestGroup"):
                test_group = tag
        comment = test['test_comment']
        steps = []
        for s in test['trace']['steps']:
            steps.append(ArchitectureRequestRecord.ArchStep(s['name'], s['start_ms'], s['end_ms']))
        return ArchitectureRequestRecord( arch, response_len, start, end, elapsed, tags, test_group, comment, steps)

    @classmethod
    def load_all(cls, reload=False) -> None:
        """
        Load all the traces from json trace log
        """
        if cls.all is None or reload:
            records = []
            test_traces = Architecture.get_trace_records()
            for trace in test_traces:
                records.append(ArchitectureRequestRecord.from_dict(trace))

            cls.all = records


class TestGroup:
    all: Dict[str, TestGroup] = None

    def __init__(self, test_group:str):
        self.arch_request_records: List[ArchitectureRequestRecord] = []
        self.test_group = test_group
        self.comment = None
        self.start = None
        self.end = None
        self.elapsed = None
        self.architectures = set()

    @property
    def num_archs(self) -> int:
        return len(self.architectures)

    @property
    def num_tests(self) -> int:
        return len(self.arch_request_records)

    @property
    def num_tests_per_arch(self) -> int:
        # Should always be an even number but cast to int just in case
        return int(self.num_tests / self.num_archs)

    def arch_request_records_by_arch(self) -> Dict[List[ArchitectureRequestRecord]]:
        grouped = {}
        for arr in self.arch_request_records:
            if arr.arch not in grouped:
                grouped[arr.arch] = []
            grouped[arr.arch].append(arr)
        return grouped

    def summary_stats_by_arch(self) -> List[Dict]:
        arch_records = self.arch_request_records_by_arch()
        arch_names = list(arch_records.keys())
        arch_names.sort()
        stats = []
        for a in arch_names:
            stat_pack = {'arch_name': a, 'elapsed': [rec.elapsed for rec in arch_records[a]],
                         'response_len': [rec.response_len for rec in arch_records[a]], 'steps': []}
            for i in range(len(arch_records[a][0].steps)):
                stat_pack['steps'].append({'step_name': arch_records[a][0].steps[i].name})
                num_recs = len(arch_records[a])
                total_elapsed = 0
                for j in range(num_recs):
                    total_elapsed += arch_records[a][j].steps[i].elapsed
                stat_pack['steps'][-1]['mean_elapsed'] = total_elapsed / num_recs
            stats.append(stat_pack)
        return stats

    def add_record(self, arr: ArchitectureRequestRecord) -> None:
        if arr.test_group != self.test_group:
            raise ValueError("Attempted to group a test record into the wrong group")
        self.arch_request_records.append(arr)
        self.architectures.add(arr.arch)
        if self.comment is None:
            self.comment = arr.comment
        if self.start is None or self.start > arr.start:
            self.start = arr.start
        if self.end is None or self.end < arr.end:
            self.end = arr.end
        self.elapsed = self.end - self.start

    @classmethod
    def load_db_records(cls) -> List[TestGroup]:
        db_file = os.path.join(data_dir, 'sqlite', 'test_records.db')
        con = sqlite3.connect(db_file)
        cur = con.cursor()

        sql = "SELECT id, test_group from test_groups"
        cur.execute(sql)
        tg_id_names =  [(r[0], r[1]) for r in cur.fetchall()]
        test_groups: List[TestGroup] = []

        for tg_id, tg_name in tg_id_names:
            tg = TestGroup(tg_name)
            sql = f"SELECT id, arch_name, response_len, start, end, comment FROM arch_requests WHERE test_group_id={tg_id}"
            cur.execute(sql)
            arch_requests = [(r[0], r[1], r[2], r[3], r[4], r[5]) for r in cur.fetchall()]
            for ar_id, ar_arch_name, ar_resp_len, ar_start, ar_end, ar_comment in arch_requests:
                sql = f"SELECT name, start, end FROM arch_req_steps WHERE arch_req_id={ar_id}"
                cur.execute(sql)
                steps = [ArchitectureRequestRecord.ArchStep(r[0], r[1], r[2]) for r in cur.fetchall()]
                arch_req_record = ArchitectureRequestRecord(ar_arch_name, ar_resp_len, ar_start, ar_end, (ar_end - ar_start), [tg_name], tg_name, f"(DB persisted) {ar_comment}", steps)
                tg.add_record(arch_req_record)
            test_groups.append(tg)
        return test_groups

    @classmethod
    def load_all(cls, reload=False):
        if cls.all is None or reload:
            ArchitectureRequestRecord.load_all(reload=reload)
            records = {}
            for arr in ArchitectureRequestRecord.all:
                if arr.test_group is not None:
                    if arr.test_group not in records:
                        records[arr.test_group] = TestGroup(arr.test_group)
                    records[arr.test_group].add_record(arr)
            db_records = cls.load_db_records()
            for test_group in db_records:
                records[test_group.test_group] = test_group
            cls.all = records

    @classmethod
    def for_test_group_tag(cls, test_group_tag: str) -> TestGroup:
        cls.load_all()
        return cls.all[test_group_tag]


def move_test_records_to_db(hf_hub_token: str) -> None:
    """
    This is an offline utility to move the transaction logs from
    a flat file to a database. To keep things simpler transaction logs
    are initially stored into a json file as a Hugging Face dataset, but
    this can get cumbersome, so this utility will move the records
    into an sqlite database.  It can be run periodically just to move
    them across.
    """
    def download_latest_json_file(hf_hub_token: str) -> Repository:
        if os.path.exists(Architecture.trace_dir):
            shutil.rmtree(Architecture.trace_dir)
        return Repository(local_dir=Architecture.trace_dir, clone_from=Architecture.save_repo_url, token=hf_hub_token)

    def create_local_db():
        db_file = os.path.join(data_dir, 'sqlite', 'test_records.db')
        con = sqlite3.connect(db_file)

        sql = "CREATE TABLE test_groups (id INTEGER PRIMARY KEY AUTOINCREMENT, test_group TEXT NOT NULL, start INTEGER NOT NULL, end INTEGER NOT NULL);"
        con.execute(sql)

        sql = "CREATE TABLE arch_requests (id INTEGER PRIMARY KEY AUTOINCREMENT, arch_name TEXT NOT NULL, response_len INTEGER NOT NULL, start INTEGER NOT NULL, end INTEGER NOT NULL, comment TEXT NOT NULL, test_group_id INTEGER NOT NULL, FOREIGN KEY (test_group_id) REFERENCES test_groups (id))"
        con.execute(sql)

        sql = "CREATE TABLE arch_req_steps (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, start INTEGER NOT NULL, end INTEGER NOT NULL, arch_req_id INTEGER NOT NULL, FOREIGN KEY (arch_req_id) REFERENCES arch_requests(id))"
        con.execute(sql)

    def get_local_db() -> sqlite3.Connection:
        db_file = os.path.join(data_dir, 'sqlite', 'test_records.db')
        if not os.path.exists(db_file):
            create_local_db()
        return sqlite3.connect(db_file)

    def load_test_group_to_db(test_group: TestGroup, con: sqlite3.Connection) -> None:
        cur = con.cursor()
        sql = f'SELECT count(*) from test_groups where test_group ="{test_group.test_group}"'
        cur.execute(sql)
        tg_not_in_db = cur.fetchall()[0][0] == 0
        if tg_not_in_db:
            sql = f'INSERT into test_groups (test_group, start, end) VALUES ("{test_group.test_group}", {test_group.start}, {test_group.end})'
            tg_id = con.execute(sql).lastrowid

            for arr in test_group.arch_request_records:
                sql = f'INSERT INTO arch_requests (arch_name, response_len, start, end, comment, test_group_id) VALUES ("{arr.arch}", {arr.response_len}, {arr.start}, {arr.end}, "{arr.comment}", {tg_id})'
                arr_id = con.execute(sql).lastrowid

                for s in arr.steps:
                    sql= f'INSERT INTO arch_req_steps (name, start, end, arch_req_id) VALUES ("{s.name}", {s.start}, {s.end}, {arr_id})'
                    con.execute(sql)
            con.commit()
        else:
            print(f"Warning TestGroup {tg_not_in_db} was not added to the DB as it already existed there")

    def load_all_test_groups_to_db(con: sqlite3.Connection) -> None:
        TestGroup.load_all()
        for tg in TestGroup.all.values():
            load_test_group_to_db(tg, con)

    # Main control flow using utility local functions above for better structure
    download_latest_json_file(hf_hub_token)
    conn = get_local_db()
    load_all_test_groups_to_db(conn)
    Architecture.wipe_trace(hf_hub_token)
    print("REMINDER: need to commit the local sqlite file to make it available to the server")


if __name__ == "__main__":
    # Expected to only be directly called for the json to db transfer - arg should be the HF token
    hf_token = sys.argv[1]
    move_test_records_to_db(hf_token)