{"cells":[{"metadata":{},"cell_type":"markdown","source":"Below is a python class I developed to assist with managing machine learning features locally on my PC.\n\nThe class is a key value store with a little extra magic.\n\nIt will automatically use disk resources when the memory limit has been hit.\n\nThe features will be stored in a sqllite backed data file on each file system resource that is added during configuration.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# 3 Example use cases\n\n    # ###############################################################################################\n    # simple usage - no persistance here\n    # ###############################################################################################\n    with TierKV(name='storename') as store1:\n        store1.set('key1', data1)\n        store1.set('key2', data2)\n        x = store1.get('key1')\n        y = store1.get('key2')\n        z = store1.get('key3') # this will return None\n        l = store1.get(['key1','key2','key3'])\n        print(x,y,z)\n        print('l is', l)\n\n    # ###############################################################################################\n    # simple usage - no persistance here\n    # ###############################################################################################\n    with TierKV(name='sample_store') as store2:\n        store2.set(['BLUE','1'], bytearray('BLUE 1', 'utf-8'))\n        store2.set(['BLUE','2'], bytearray('BLUE 2', 'utf-8'))\n        store2.set(['BLUE','3'], bytearray('BLUE 3', 'utf-8'))\n        store2.set(['BLUE','4'], bytearray('BLUE 4', 'utf-8'))\n        store2.set(['BLUE','5'], bytearray('BLUE 5', 'utf-8'))\n        store2.set(['BLUE','6'], bytearray('BLUE 6', 'utf-8'))\n        store2.set(['BLUE','7'], bytearray('BLUE 7', 'utf-8'))\n        store2.set(['RED','1'], bytearray('RED 1', 'utf-8'))\n        \n        # get all data tagged with BLUE\n        blues = store2.get('BLUE')\n        print('blues',blues)\n\n        # get 3 random data tagged with BLUE\n        blues_3 = store2.sample('BLUE',30000) # this will return all\n        blues_3 = store2.sample('BLUE',3)\n        print('blues 3', blues_3)\n\n        # get all data tagged with RED\n        reds = store2.get('RED')\n        print('reds',reds)\n        \n    # ###############################################################################################\n    # more complex use case - with persistance via filesystem\n    # configure the tiers of storage resources to utilise\n    # all limits are 'soft' limits to target, no guarantees.\n    # ###############################################################################################\n    store3 = TierKV(name='mystore', resources=[\n        # Tier 0 - use memory for speed, and limit to 90% available memory\n        ResourceMemory(limit = 0.90, compression=False),   \n        # Tier 1 - C: is a fast M2 drive, so use this next, enable compression on this resource and limit to 90% disk\n        ResourceFilesystem(path = 'c:/store/', limit = 0.90, compression=True), \n        # Tier 2 - D: is a slower SSD, so use this last, enable compression and and limit to 90% disk\n        ResourceFilesystem(path = 'd:/store/', limit = 0.90, compression=True)\n        ])\n\n    with store3:\n        store3.set(['BLUE','1'], bytearray('BLUE 1', 'utf-8'))\n        store3.set(['BLUE','2'], bytearray('BLUE 2', 'utf-8'))\n        store3.set(['BLUE','3'], bytearray('BLUE 3', 'utf-8'))\n        store3.set(['BLUE','4'], bytearray('BLUE 4', 'utf-8'))\n        store3.set(['BLUE','5'], bytearray('BLUE 5', 'utf-8'))\n        store3.set(['BLUE','6'], bytearray('BLUE 6', 'utf-8'))\n        store3.set(['BLUE','7'], bytearray('BLUE 7', 'utf-8'))\n        store3.set(['RED','1'], bytearray('RED 1', 'utf-8'))\n        \n        # get all data tagged with BLUE\n        blues = store3.get('BLUE')\n        print('blues',blues)\n\n        # get 3 random data tagged with BLUE\n        blues_3 = store3.sample('BLUE',30000) # this will return all\n        blues_3 = store3.sample('BLUE',3)\n        print('blues 3', blues_3)\n\n        # get all data tagged with RED\n        reds = store3.get('RED')\n        print('reds',reds)      ","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import hashlib\nimport os\nimport blosc\nimport psutil\nimport threading\nimport sqlite3\nimport time\nimport logging\nimport random\nimport pickle\n\n\nclass Resource():\n\n    def __init__(self, limit = 0.9, compression = False):\n        self.limit = limit\n        self.compression = compression  \n        \n    def on_exit(self):\n        pass\n    \n    def on_enter(self):\n        pass\n\n    def set_name(self, name):\n        self.name = name\n\n    def get_limit(self):\n        return self.limit\n\n\nclass ResourceMemory(Resource):\n\n    def __init__(self, limit = 0.9, compression = False):\n        super().__init__(limit, compression)\n        \n    def initialise(self):\n        self.store = {}\n        self.update_percent_used();\n\n    def update_percent_used(self):\n        self._percent_used = psutil.virtual_memory().percent/100.0\n\n    def percent_used(self):\n        return self._percent_used\n\n    def keys(self):\n        result = list(self.store.keys())\n        return result\n\n    def has_key(self, key):\n        return key in self.store\n\n    def delete_key(self, key):\n        self.store.pop(key, None)\n\n    # ###########################################################################\n    # PUBLIC API\n    # ###########################################################################\n    def set(self, key, value, bypass=False):\n        if self.compression and not bypass:\n            value = blosc.compress(value, typesize=4, cname='lz4')\n        self.store[key] = value\n\n    def get(self, key, bypass=False):\n        result = None\n        try:\n            if self.compression and not bypass:\n                result = blosc.decompress(self.store[key])\n            else:\n                result = self.store[key]\n        except:\n            pass\n        return result\n\n\nclass ResourceFilesystem(Resource):\n\n    def __init__(self, path, limit = 0.90, compression = False):\n        super().__init__(limit, compression)\n        self.path = path\n\n    def initialise(self):\n        directory = os.path.join(self.path, self.name)\n        if not os.path.exists(directory):\n            os.makedirs(directory)\n        filename = os.path.join(self.path, self.name, 'TierKV.db')\n        first_time = not os.path.exists(filename)\n        self.sqliteConnection = sqlite3.connect(filename, check_same_thread = False)\n        if first_time:\n            self.cursor = self.sqliteConnection.cursor()\n            self.cursor.execute('CREATE TABLE IF NOT EXISTS data (k TEXT PRIMARY KEY, v BLOB);')\n            # self.cursor.execute('PRAGMA auto_vacuum = INCREMENTAL;')\n            self.cursor.execute('PRAGMA auto_vacuum = NONE;')\n            self.cursor.execute('VACUUM;')\n            self.cursor.close()\n        self.update_percent_used();\n\n    def on_enter(self):\n        self.cursor = self.sqliteConnection.cursor()\n\n    def on_exit(self):\n        self.sqliteConnection.commit()\n        self.cursor.close()\n\n    def update_percent_used(self):\n        self._percent_used = psutil.disk_usage(self.path).percent/100.0\n\n    def percent_used(self):\n        return self._percent_used\n\n    def keys(self):\n        self.cursor.execute('SELECT k FROM data')\n        result = []\n        for row in self.cursor.fetchall():\n            result.append(row[0])\n        return result\n\n    def has_key(self, key):\n        self.cursor.execute('SELECT k FROM data WHERE k=?', (key,))\n        row = self.cursor.fetchone()\n        result = row is not None\n        return result\n\n    def delete_key(self, key):\n        # if self.has_key(key) is not None:\n        self.cursor.execute('DELETE FROM data WHERE k=?', (key,))\n\n    # ###########################################################################\n    # PUBLIC API\n    # ###########################################################################\n    def set(self, key, value, bypass=False):\n        if self.compression and not bypass:\n            value = blosc.compress(value, typesize=4, cname='lz4')\n        self.cursor.execute('INSERT INTO data (k,v) VALUES (?,?) ON CONFLICT(k) DO UPDATE SET v=?;', (key,value,value))\n\n    def get(self, key, bypass=False):\n        self.cursor.execute('SELECT v FROM data WHERE k=?', (key,))\n        row = self.cursor.fetchone()\n        result = None\n        try:\n            if row is not None:\n                if self.compression and not bypass:\n                    result = blosc.decompress(row[0])\n                else:\n                    result = row[0]\n        except:\n            pass\n        return result\n\n\nclass TierKV():\n\n    def __init__(self, name = 'default', resources=[ResourceMemory(limit = 0.90)], cache=False):\n        self.tag_keys_map =  {} # tag_keys_map  tag  --> set(keys)\n        self.key_tier_map =  {} # key_tier_map  key  --> tier\n        self.tier_keys_map = {} # tier_keys_map tier --> set(keys)\n\n        self.cache = cache\n        self.name = name\n        self.resources = resources        \n        self.lock = threading.RLock()\n\n        self.lock.acquire()\n        for resource in self.resources:\n            resource.set_name(self.name)\n        for tier, resource in enumerate(self.resources):\n            resource.initialise()\n            self.tier_keys_map[tier] = set()\n        self.lock.release()\n\n    def __enter__(self):\n        for resource in self.resources:\n            resource.on_enter()\n        self.load()\n        self.thread_run = True\n        self.thread = threading.Thread(target=self.__worker_thread)\n        self.thread.start()\n        return self\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        self.thread_run = False\n        self.thread.join()\n        self.save()\n\n        # do any final commit on resource\n        for resource in self.resources:\n            resource.on_exit()\n\n    def __worker_thread(self):\n        counter=0\n        while(self.thread_run):\n            time.sleep(0)\n            counter = counter+1\n            if (counter % 10000) == 0:\n                self.lock.acquire()\n                for resource in self.resources:\n                    resource.update_percent_used()\n                self.__rebalance_tiers()\n                self.lock.release()\n\n    def __find_free_tier_above(self, tier):\n        self.lock.acquire()\n        for t in range(tier-1, -1, -1):\n            if self.resources[t].percent_used() < self.resources[t].get_limit():\n                self.lock.release()\n                return t\n        self.lock.release()\n        return None\n\n    def __find_free_tier_below(self, tier):\n        self.lock.acquire()\n        for t in range(tier+1, len(self.resources)):\n            if self.resources[t].percent_used() < self.resources[t].get_limit():\n                self.lock.release()\n                return t\n        self.lock.release()\n        return None\n\n    def __key(self, tags):\n        return hashlib.blake2s(bytes(''.join(tags), encoding='utf8'), digest_size=20).digest()\n\n    def __move_up(self, key, tier, data):\n        self.lock.acquire()\n        above = self.__find_free_tier_above(tier)\n        if above is not None:\n            bypass=False\n            if self.resources[tier].compression == self.resources[above].compression:\n                bypass=True\n            r_src = self.resources[tier]\n            r_dst = self.resources[above]\n            r_dst.set(key, data, bypass)\n            # TODO: I wonder if this is ok?  r_src.delete_key(key)\n            self.key_tier_map[key] = above\n            self.tier_keys_map[tier].remove(key)\n            self.tier_keys_map[above].add(key)\n        self.lock.release()\n\n    def __move_down(self, key, tier):\n        self.lock.acquire()    \n        below = self.__find_free_tier_below(tier)\n        if below is not None:\n            bypass=False\n            if self.resources[tier].compression == self.resources[below].compression:\n                bypass=True\n            r_src = self.resources[tier]\n            r_dst = self.resources[below]\n            r_dst.set(key, r_src.get(key, bypass), bypass)\n            r_src.delete_key(key)\n            self.key_tier_map[key] = below\n            self.tier_keys_map[tier].remove(key)\n            self.tier_keys_map[below].add(key)\n        self.lock.release()\n\n    def __rebalance_tiers(self):\n        self.lock.acquire()\n        if len(self.resources)>1:\n            # only tier 0\n            tier = 0\n            # slowly migrate 1 record at a time to free up 5% head room\n            if len(self.tier_keys_map[tier]) > 1:\n                if (self.resources[tier].percent_used() > (self.resources[tier].get_limit()-0.05)) and (self.resources[tier+1].percent_used() < self.resources[tier+1].get_limit()):\n                    # evict a random key\n                    for k in random.sample(self.tier_keys_map[tier], 1):  \n                        self.__move_down(k, tier)\n        self.lock.release()\n\n    def set(self, tags, data):\n\n        # ensure we have a list\n        if not isinstance(tags, list):\n            tags = [tags]\n\n        # calculate the key\n        key = self.__key(tags)\n\n        self.lock.acquire()\n\n        # first remove from existing tier\n        if key in self.key_tier_map:\n            tier = self.key_tier_map[key]\n            self.resources[tier].delete_key(key)\n            self.tier_keys_map[tier].remove(key)\n\n        # add into first tier that has space left\n        for tier, resource in enumerate(self.resources):\n            if resource.percent_used() < resource.get_limit():\n                resource.set(key, data)\n                self.key_tier_map[key] = tier\n                self.tier_keys_map[tier].add(key)\n\n                for tag in tags:\n                    if not tag in self.tag_keys_map:\n                        self.tag_keys_map[tag] = set()\n                    self.tag_keys_map[tag].add(key)\n\n                self.lock.release()\n                return key\n\n        # only get here if we are out of resource space\n        self.lock.release()\n        raise Exception('out of resources')\n\n    def sample(self, tags, k):\n        \n        # ensure we have a list\n        if not isinstance(tags, list):\n            tags = [tags]\n\n        self.lock.acquire()\n        result = []\n        seen = set()\n        for tag in tags:\n            if tag in self.tag_keys_map:\n                k =  min(k, len(self.tag_keys_map[tag]))\n                for key in random.sample(self.tag_keys_map[tag], k):\n                    if key not in seen:\n                        tier = self.key_tier_map[key]\n                        data = self.resources[tier].get(key)\n                        if data is not None:\n                            result.append(data)        \n                            seen.add(key)\n                            if self.cache and (tier>0):\n                                self.__move_up(key, tier, data)\n        self.lock.release()\n\n        return result\n\n    def get(self, tags):\n        \n        # ensure we have a list\n        if not isinstance(tags, list):\n            tags = [tags]\n\n        self.lock.acquire()\n        result = []\n        seen = set()\n        for tag in tags:\n            if tag in self.tag_keys_map:\n                for key in self.tag_keys_map[tag]:\n                    if key not in seen:\n                        tier = self.key_tier_map[key]\n                        data = self.resources[tier].get(key)\n                        if data is not None:\n                            result.append(data)        \n                            seen.add(key)\n                            if self.cache and (tier>0):\n                                self.__move_up(key, tier, data)\n        self.lock.release()\n\n        return result\n\n    def save(self):\n        self.lock.acquire()\n        if len(self.resources)>1:\n\n            # move everything down from tier 0 into a persistant tier\n            tier = 0\n            for key in self.resources[tier].keys():\n                self.__move_down(key, tier)\n\n            # save all the dictionaries too\n            self.resources[1].set('tag_keys_map.pkl', pickle.dumps(self.tag_keys_map))\n            self.resources[1].set('key_tier_map.pkl', pickle.dumps(self.key_tier_map))\n            self.resources[1].set('tier_keys_map.pkl', pickle.dumps(self.tier_keys_map))\n\n        self.lock.release()\n\n    def load(self):\n        self.lock.acquire()\n        if len(self.resources)>1:\n            if self.resources[1].get('tag_keys_map.pkl') is not None:\n                self.tag_keys_map = pickle.loads(self.resources[1].get('tag_keys_map.pkl'))\n                self.key_tier_map = pickle.loads(self.resources[1].get('key_tier_map.pkl'))\n                self.tier_keys_map = pickle.loads(self.resources[1].get('tier_keys_map.pkl'))\n        self.lock.release()\n\n    def has_tags(self, tags):\n        self.lock.acquire()\n        result = False\n        for tag in tags:\n            if tag in self.tag_keys_map.keys():\n                result = True\n                break\n        self.lock.release()\n        return result\n\n    def keys(self):\n        return list(self.key_tier_map.keys())\n\n    def get_by_key(self, key):      \n        self.lock.acquire()\n        result = []\n        tier = self.key_tier_map[key]\n        data = self.resources[tier].get(key)\n        if data is not None:\n            result.append(data)        \n            if self.cache and (tier>0):\n                self.__move_up(key, tier, data)\n        self.lock.release()\n        return result","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}