wirehead 0.1.0__tar.gz

wirehead-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 neuroneural/wirehead
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

wirehead-0.1.0/PKG-INFO

@@ -0,0 +1,4 @@
+Metadata-Version: 2.1
+Name: wirehead
+Version: 0.1.0
+License-File: LICENSE

wirehead-0.1.0/README.md

@@ -0,0 +1,201 @@
+# Wirehead #
+
+Caching system for scaling of synthetic data generators using MongoDB
+
+![wirehead-diagram](assets/training-diagram.png)
+
+---
+
+# I. Installation 
+
+## 1. MongoDB Setup (For Development/Testing Only)
+
+- [Ubuntu Setup](#a-quick-mongodb-setup-ubuntu)
+- [macOS Setup](#b-quick-mongodb-setup-macos)
+
+**Important Note:** The following instructions are for development and testing purposes only. For production deployments, please refer to the [official MongoDB documentation](https://www.mongodb.com/docs/manual/administration/install-community/) for secure and proper installation guidelines.
+
+#### a. Quick MongoDB Setup ([Ubuntu](https://www.mongodb.com/docs/manual/tutorial/install-mongodb-on-ubuntu/)):
+
+```bash
+sudo apt-get install gnupg curl
+curl -fsSL https://www.mongodb.org/static/pgp/server-7.0.asc | \
+   sudo gpg -o /usr/share/keyrings/mongodb-server-7.0.gpg \
+   --dearmor
+echo "deb [ arch=amd64,arm64 signed-by=/usr/share/keyrings/mongodb-server-7.0.gpg ] https://repo.mongodb.org/apt/ubuntu jammy/mongodb-org/7.0 multiverse" | sudo tee /etc/apt/sources.list.d/mongodb-org-7.0.list
+sudo apt-get update
+sudo apt-get install -y mongodb-org
+```
+
+```bash
+# Run MongoDB
+sudo systemctl start mongod
+```
+
+```bash
+# Stop MongoDB
+sudo systemctl stop mongod
+```
+
+#### b. Quick MongoDB Setup ([MacOS](https://www.mongodb.com/docs/manual/tutorial/install-mongodb-on-os-x/)):
+
+```bash
+brew tap mongodb/brew
+brew update
+brew install mongodb-community@7.0
+```
+
+```bash
+# Run MongoDB
+brew services start mongodb-community@7.0
+```
+
+```bash
+# Run MongoDB
+brew services stop mongodb-community@7.0
+```
+
+**Warning:** These instructions are simplified for ease of setup in a development environment. For production use, ensure proper security measures, authentication, and follow best practices as outlined in the official MongoDB documentation.
+Installing and deploying MongoDB
+
+## 2. Create virtual environment
+```bash
+# Note:
+# python version doesn't necessarily have to be 3.10
+# but this gives better support for some generation pipelines
+
+# Conda
+conda create -n wirehead python=3.10
+conda activate wirehead
+
+# venv
+python3.10 -m venv wirehead 
+source venv/bin/activate
+```
+
+## 3. Install wirehead:
+```bash
+git clone git@github.com:neuroneural/wirehead.git
+cd wirehead
+pip install -e .
+```
+
+## 4. Run the test
+```bash
+cd examples/unit
+chmod +x test.sh
+./test.sh
+```
+
+# II. Usage 
+
+See examples/unit for a minimal example 
+
+## 1. Manager
+```python
+from wirehead import WireheadManager
+
+if __name__ == "__main__":
+    wirehead_runtime = WireheadManager(config_path="config.yaml")
+    wirehead_runtime.run_manager()
+```
+
+## 2. Generator
+
+```python
+import numpy as np
+from wirehead import WireheadGenerator 
+
+def create_generator():
+    while True: 
+        img = np.random.rand(256,256,256)
+        lab = np.random.rand(256,256,256)
+        yield (img, lab)
+
+if __name__ == "__main__":
+    brain_generator     = create_generator()
+    wirehead_runtime    = WireheadGenerator(
+        generator = brain_generator,
+        config_path = "config.yaml" 
+    )
+    wirehead_runtime.run_generator()
+```
+
+## 3. Dataset
+```python
+import torch
+from wirehead import MongoheadDataset
+
+dataset = MongoheadDataset(config_path = "config.yaml")
+
+idx = [0] 
+data = dataset[idx]
+sample, label = data[0]['input'], data[0]['label']
+```
+
+# III. Config guide
+
+All wirehead configs live inside yaml files, and must be specified when declaring wirehead manager, generator and dataset objects. For the system to work, all components must use the __same__ configs.
+
+## 1. Basic configs:
+```yaml
+MONGOHOST -- IP address or hostname for machine running MongoDB instance
+DBNAME -- MongoDB database name
+PORT -- Port for MongoDB instance. Defaults to 27017
+SWAP_CAP -- Size cap for read and write collections. bigger means bigger cache, and less frequent swaps. The total memory used by wirehead can be calculated with:
+        SWAP_CAP * SIZE OF YIELDED TUPLE * 2
+```
+
+## 2. Advanced configs:
+```yaml
+SAMPLE -- Array of strings denoting name of samples in data tuple. 
+WRITE_COLLECTION   -- Name of write collection (generators push to this)
+READ_COLLECTION    -- Name of read colletion (dataset reads from this)
+COUNTER_COLLECTION -- Name of counter collection for manager metrics
+TEMP_COLLECTION    -- Name of temporary collection used for moving data during swap
+CHUNKSIZE          -- Number of megabytes used for chunking data
+```
+
+---
+
+# IV. Generator example
+
+See a simple example in [examples/unit/generator.py](examples/unit/generator.py) or a Synthseg example in [examples/synthseg/generator.py](examples/synthseg/generator.py)
+
+Wirehead's [WireheadGenerator](https://github.com/neuroneural/wirehead/blob/main/wirehead/generator.py) object takes in a generator, which is a python generator function. This function yields a tuple containing numpy arrays. The number of samples in this tuple should match the number of strings specified in SAMPLE in config.yaml
+
+## 1. Set SAMPLE in "config.yaml" (note the number of keys)
+```yaml
+SAMPLE: ["a", "b"]
+```
+
+## 2. Create a generator function, which yields the same number of objects
+```python
+def create_generator():
+    while True: 
+        a = np.random.rand(256,256,256)
+        b = np.random.rand(256,256,256)
+        yield (a, b)
+```
+
+## 3. Insert config file path and generator function into WireheadGenerator
+```python
+generator = create_generator()
+runtime = WireheadGenerator(
+    generator = generator,
+    config_path = "config.yaml" 
+)
+```
+
+## 4. Press play
+```python
+runtime.run_generator() # runs an infinite loop
+```
+
+---
+
+# Citation/Contact
+
+This code is under [MIT](https://github.com/neuroneural/wirehead/blob/main/LICENSE) licensing
+
+If you have any questions specific to the Wirehead pipeline, please raise an issue or contact us at mdoan4@gsu.edu

wirehead-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

wirehead-0.1.0/setup.py

@@ -0,0 +1,21 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='wirehead',
+    version='0.1.0',
+    packages=find_packages(),
+    install_requires=[
+        'pymongo',
+        'torch',
+        'numpy',
+        'PyYaml',
+        # List your package dependencies here
+        # Example: 'numpy>=1.19.0',
+    ],
+    entry_points={
+        'console_scripts': [
+            # Define any command-line entry points here
+            # Example: 'your-command = your_package.module:function',
+        ],
+    },
+)

wirehead-0.1.0/wirehead/__init__.py

@@ -0,0 +1,4 @@
+from wirehead.dataset import MongoheadDataset
+from wirehead.dataset import MongoTupleheadDataset
+from wirehead.manager import WireheadManager
+from wirehead.generator import WireheadGenerator

wirehead-0.1.0/wirehead/dataset.py

@@ -0,0 +1,253 @@
+""" Wirehead dataset class for MongoDB """
+
+import io
+import os
+import sys
+import time
+import yaml
+import torch
+from torch.utils.data import Dataset
+from pymongo import MongoClient
+from pymongo.errors import OperationFailure, OperationFailure
+
+
+def unit_interval_normalize(img):
+    """Unit interval preprocessing"""
+    img = (img - img.min()) / (img.max() - img.min())
+    return img
+
+
+def quantile_normalize(img, qmin=0.01, qmax=0.99):
+    """Unit interval preprocessing"""
+    img = (img - img.quantile(qmin)) / (img.quantile(qmax) - img.quantile(qmin))
+    return img
+
+
+def binary_to_tensor(tensor_binary):
+    """ Converts a binary io buffer to a torch tensor """
+    buffer = io.BytesIO(tensor_binary)
+    tensor = torch.load(buffer)
+    return tensor
+
+
+class MongoheadDataset(Dataset):
+    """
+    A dataset for fetching batches of records from a MongoDB
+    """
+
+    def __init__(self,
+                 config_path="",
+                 timeout=60,
+                 collection=None,
+                 sample=("data", "label"),
+                 transform=binary_to_tensor,
+                 normalize=lambda x: x,
+                 id="id",
+                 keeptrying=True):
+        """Constructor
+        :param config_path: path to wirehead config .yaml file
+        :param indices: a set of indices to be extracted from the collection
+        :param transform: a function to be applied to each extracted record
+        :param collection: pymongo collection to be used
+        :param sample: a pair of fields to be fetched as `input` and `label`
+                       , e.g. (`T1`, `label104`)
+        :param id: the field to be used as an index. The `indices` are values of this field
+        :param keeptrying: whether to attempt a refetch if first attempt fails
+        :returns: an object of MongoheadDataset class
+        """
+        self.id = id
+        self.normalize = normalize
+        self.transform = transform
+        self.keeptrying = keeptrying    # retries if fetch fails
+        self.fields = {"id": 1, "chunk": 1, "kind": 1, "chunk_id": 1}
+        self.timeout = timeout
+
+        if config_path != "" and os.path.exists(config_path):
+            self.load_from_yaml(config_path)
+
+        else:
+            self.collection = collection
+            self.sample = sample
+
+        self.indices = self.get_indeces()
+
+    def load_from_yaml(self, config_path):
+        """
+        Loads config options from config_path
+        """
+        print("Dataset: config loaded from " + config_path)
+        with open(config_path, "r") as file:
+            config = yaml.safe_load(file)
+        dbname = config.get('DBNAME')
+        mongohost = config.get('MONGOHOST')
+        port = config.get('PORT') if config.get('PORT') is not None else 27017
+        client = MongoClient("mongodb://" + mongohost + ":" + str(port))
+
+        db = client[dbname]
+        self.wait_for_data(db)
+        read_collection = config.get("READ_COLLECTION")
+        self.collection = db[read_collection]
+        self.sample = tuple(config.get("SAMPLE"))
+
+    def wait_for_data(self, db, timeout=600, check_interval=10):
+        """
+        Prevents data object from reading before data is ready.
+        Raises:
+            TimeoutError: If the wait time exceeds the timeout
+        """
+        status_collection = db["status"]
+        start_time = time.time()
+
+        def check_status():
+            try:
+                latest_status = status_collection.find_one(sort=[("_id", -1)])
+                if latest_status is None:
+                    print("Dataset: Database is empty, waiting for data...")
+                    return False
+                if not latest_status.get("swapped", False):
+                    print("Dataset: Swap has not happened, waiting...")
+                    return False
+                return True
+            except (ConnectionFailure, OperationFailure) as e:
+                print(f"Dataset: Error accessing database: {e}")
+                return False
+
+        while time.time() - start_time < timeout:
+            if check_status():
+                print("Dataset: Data is ready")
+                return
+            time.sleep(check_interval)
+
+        raise TimeoutError("Dataset: Waited too long for data to be ready")
+
+    def get_indeces(self):
+        """
+        Retrieve the index array of samples in read collection
+        """
+        last_post = self.collection['bin'].find_one(sort=[(self.id, -1)])
+
+        if last_post is None:
+            print("Empty collection, exiting")
+            sys.exit()
+        num_examples = int(last_post[self.id] + 1)
+        return range(num_examples)
+
+    def __len__(self):
+        return len(self.indices)
+
+    def make_serial(self, samples_for_id, kind):
+        """
+        Converts collection chunks into a contiguous byte sequence
+        """
+        return b"".join([
+            sample["chunk"] for sample in sorted(
+                (sample for sample in samples_for_id if sample["kind"] == kind),
+                key=lambda x: x["chunk_id"],
+            )
+        ])
+
+    def retry_on_eof_error(retry_count, verbose=False):
+        """
+        Error handling for reads that happen mid swap
+        """
+
+        def decorator(func):
+
+            def wrapper(self, batch, *args, **kwargs):
+                myException = Exception    # Default Exception if not overwritten
+                for attempt in range(retry_count):
+                    try:
+                        return func(self, batch, *args, **kwargs)
+                    except (
+                            EOFError,
+                            OperationFailure,
+                    ) as exception:    # Specifically catching EOFError
+                        if self.keeptrying:
+                            if verbose:
+                                print(
+                                    f"EOFError caught. Retrying {attempt+1}/{retry_count}"
+                                )
+                            time.sleep(1)
+                            continue
+                        else:
+                            raise exception
+                raise myException("Failed after multiple retries.")
+
+            return wrapper
+
+        return decorator
+
+    @retry_on_eof_error(retry_count=3, verbose=True)
+    def __getitem__(self, batch):
+        """
+        Fetch all samples for ids in the batch and where 'kind' is either
+        data or label as specified by the sample parameter
+        """
+        samples = list(self.collection["bin"].find(
+            {
+                self.id: {
+                    "$in": [self.indices[_] for _ in batch]
+                },
+                "kind": {
+                    "$in": self.sample
+                },
+            },
+            self.fields,
+        ))
+        results = {}
+        for id in batch:
+            # Separate samples for this id
+            samples_for_id = [
+                sample for sample in samples if sample[self.id] == self.indices[id]
+            ]
+
+            # Separate processing for each 'kind'
+            data = self.make_serial(samples_for_id, self.sample[0])
+            label = self.make_serial(samples_for_id, self.sample[1])
+
+            # Add to results
+            results[id] = {
+                "input": self.normalize(self.transform(data).float()),
+                "label": self.transform(label),
+            }
+        return results
+
+class MongoTupleheadDataset(MongoheadDataset):
+    """
+    A dataset for fetching batches of records from a MongoDB
+    Returns a tuple instead of a dict
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @MongoheadDataset.retry_on_eof_error(retry_count=3, verbose=True)
+    def __getitem__(self, batch):
+        """
+        Fetch all samples for ids in the batch and return as tuples
+        """
+        samples = list(self.collection["bin"].find(
+            {
+                self.id: {
+                    "$in": [self.indices[_] for _ in batch]
+                },
+                "kind": {
+                    "$in": self.sample
+                },
+            },
+            self.fields,
+        ))
+        results = []
+        for id in batch:
+            samples_for_id = [
+                sample for sample in samples if sample[self.id] == self.indices[id]
+            ]
+
+            data = self.make_serial(samples_for_id, self.sample[0])
+            label = self.make_serial(samples_for_id, self.sample[1])
+
+            data = self.normalize(self.transform(data).float())
+            label = self.transform(label)
+
+            results.append((data, label))
+        return results

wirehead-0.1.0/wirehead/generator.py

@@ -0,0 +1,127 @@
+""" Wirehead Generator Class """
+
+import io
+import os
+import time
+import yaml
+import bson
+import torch
+from pymongo import MongoClient, ReturnDocument
+
+
+class WireheadGenerator():
+    """
+    Wirehead runtime class, which wraps around the generator
+    and manager runtimes.
+    """
+
+    def __init__(self, generator, config_path, n_samples = 1000):
+        if config_path is None or os.path.exists(config_path) is False:
+            print("No valid config specified, exiting")
+            return
+        self.load_from_yaml(config_path)
+        self.generator = generator
+        self.n_samples = n_samples
+
+    def load_from_yaml(self, config_path):
+        """ Loads manager configs from config_path """
+        with open(config_path, 'r', encoding='utf-8') as file:
+            config = yaml.safe_load(file)
+        dbname = config.get('DBNAME')
+        mongohost = config.get('MONGOHOST')
+        port = config.get('PORT') if config.get('PORT') is not None else 27017
+        client = MongoClient("mongodb://" + mongohost + ":" + str(port))
+
+        self.db = client[dbname]
+        self.swap_cap = config.get('SWAP_CAP')
+        self.sample = tuple(config.get("SAMPLE"))
+        self.chunksize = config.get("CHUNKSIZE")
+        self.collectionw = config.get("WRITE_COLLECTION") + ".bin"
+        self.collectionc = config.get("COUNTER_COLLECTION")
+
+    def chunkify(self, data, index):
+        """
+        Converts a tuple of tensors and their labels into
+        a list of chunks of serialized objects for mongodb
+        """
+
+        def chunk_binobj(tensor_compressed, idx, kind, chunksize):
+            """
+            Convert chunksize from megabytes to bytes
+            """
+            chunksize_bytes = chunksize * 1024 * 1024
+            # Calculate the number of chunks
+            num_chunks = len(tensor_compressed) // chunksize_bytes
+            if len(tensor_compressed) % chunksize_bytes != 0:
+                num_chunks += 1
+            # Yield chunks
+            for i in range(num_chunks):
+                start = i * chunksize_bytes
+                end = min((i + 1) * chunksize_bytes, len(tensor_compressed))
+                chunk = tensor_compressed[start:end]
+                yield {
+                    "id": idx,
+                    "chunk_id": i,
+                    "kind": kind,
+                    "chunk": bson.Binary(chunk),
+                }
+
+        def tensor2bin(tensor):
+            """
+            Seralize a torch tensor into an IO buffer
+            """
+            tensor_1d = tensor.to(torch.uint8)
+            buffer = io.BytesIO()
+            torch.save(tensor_1d, buffer)
+            tensor_binary = buffer.getvalue()
+            return tensor_binary
+
+        chunks = []
+        binobj = data
+        kinds = self.sample
+        for i, kind in enumerate(kinds):
+            chunks += list(
+                chunk_binobj(tensor2bin(torch.from_numpy(binobj[i])), index, kind,
+                             self.chunksize))
+        return chunks
+
+    def push_chunks(self, chunks):
+        """ Pushes chunkified tensors to mongodb, with error handling"""
+        collection_bin = self.db[self.collectionw]
+        try:
+            collection_bin.insert_many(chunks)
+        except Exception as exception:
+            print(f"Generator: An error occurred: {exception}, are you swapping?")
+            time.sleep(1)
+
+    def get_current_idx(self):
+        """ Get current index of sample in write collection """
+        dbc = self.db[self.collectionc]
+        counter_doc = dbc.find_one_and_update(
+            {"_id": "uniqueFieldCounter"},
+            {"$inc": {
+                "sequence_value": 1
+            }},
+            return_document=ReturnDocument.BEFORE,
+        )
+        return counter_doc["sequence_value"]
+
+    def generate_and_insert(self):
+        """ Fetch from generator and inserts into mongodb """
+        # 0. Fetch data from generator
+        data = next(self.generator)
+        # 1. Get the correct index for this current sample
+        index = self.get_current_idx()
+        # 2. Turn the data into a list of serialized chunks
+        chunks = self.chunkify(data, index)
+        # 3. Push to mongodb + error handling
+        if index < self.swap_cap:
+            self.push_chunks(chunks)
+
+    def run_generator(self):
+        """ Initializes and runs a SynthSeg brain generator in a loop,
+            preprocesses, then pushes to mongoDB"""
+        print("Generator: Initialized")
+        n_samples = self.n_samples
+        for _ in range(n_samples):
+            self.generate_and_insert()

wirehead-0.1.0/wirehead/manager.py

@@ -0,0 +1,117 @@
+""" Wirehead Manager Class """
+
+import os
+import time
+import yaml
+from pymongo import MongoClient, ASCENDING
+
+
+class WireheadManager():
+    """
+    Manages the state of the mongo collections in Wirehead.
+
+    :param config_path: path to yaml file containing wirehead configs
+    """
+
+    def __init__(self, config_path):
+        if config_path is None or os.path.exists(config_path) is False:
+            print("No valid config specified, exiting")
+            return
+        self.load_from_yaml(config_path)
+
+    def load_from_yaml(self, config_path):
+        """
+        Loads manager configs from config_path.
+        """
+        with open(config_path, 'r', encoding='utf-8') as file:
+            config = yaml.safe_load(file)
+
+        dbname = config.get('DBNAME')
+        mongohost = config.get('MONGOHOST')
+        port = config.get('PORT') if config.get('PORT') is not None else 27017
+        client = MongoClient("mongodb://" + mongohost + ":" + str(port))
+
+        self.db = client[dbname]
+        self.swap_cap = config.get('SWAP_CAP')
+        self.collectionw = config.get("WRITE_COLLECTION") + ".bin"
+        self.collectionr = config.get("READ_COLLECTION") + ".bin"
+        self.collectionc = config.get("COUNTER_COLLECTION")
+        self.collectiont = config.get("TEMP_COLLECTION") + ".bin"
+
+    def run_manager(self):
+        """
+        Initializes the database manager, swaps and cleans the database whenever swap_cap is hit.
+        """
+        print("Manager: Initialized")
+        self.db["status"].insert_one({"swapped": False})
+        self.reset_counter_and_collection()
+        generated = 0
+        while True:
+            generated = self.watch_and_swap(generated)
+
+    def verify_collection_integrity(self, collection):
+        """
+        Verifies collection contains contiguous elements with id 0..swap_cap
+        """
+        unique_ids_count = len(collection.distinct("id"))
+        assert (
+            unique_ids_count == self.swap_cap
+        ), f"Manager: Expected {self.swap_cap} unique ids, found {unique_ids_count}"
+        expected_ids_set = set(range(self.swap_cap))
+        actual_ids_set = set(collection.distinct("id"))
+        assert (expected_ids_set == actual_ids_set
+                ), "Manager: The ids aren't continuous from 0 to self.swap_cap - 1"
+
+    def reset_counter_and_collection(self):
+        """
+        Delete all documents in the main collection that have creeped in
+        between the renaming and now. This operation is within a transaction.
+        """
+        dbw = self.db[self.collectionw]
+        dbc = self.db[self.collectionc]
+        dbw.delete_many({})    # wipe the write collection
+        # Reset the counter to zero
+        _result = dbc.update_one(
+            {"_id": "uniqueFieldCounter"},    # Query part: the document to match
+            {"$set": {
+                "sequence_value": 0
+            }},    # Update part: what to set if the document is matched/found
+            upsert=True,
+        )
+        dbw.delete_many({})
+        dbw.create_index([("id", ASCENDING)], background=True)
+
+    def swap(self, generated):
+        """
+        Moves data from write collection to read collection
+        Deletes old write collection
+        Maintains data integrity in between
+        """
+        time.sleep(2)    # Buffer for incomplete ops
+        generated += self.swap_cap
+        print("\n----swap----")
+        print(f"Manager: Generated samples so far {generated}")
+        self.db[self.collectionw].rename(self.collectiont, dropTarget=True)
+        # Now atomically reset the counter to 0 and delete whatever records
+        # may have been written between the execution of the previous line
+        # and the next
+        self.reset_counter_and_collection()    # this is atomic
+        result = self.db[self.collectiont].delete_many(
+            {"id": {
+                "$gt": self.swap_cap - 1
+            }})
+        # Print the result of the deletion
+        print(f"Manager: Documents deleted: {result.deleted_count}")
+        self.verify_collection_integrity(self.db[self.collectiont])
+        self.db[self.collectiont].rename(self.collectionr, dropTarget=True)
+        self.db["status"].insert_one({"swapped": True})
+        return generated
+
+    def watch_and_swap(self, generated):
+        """
+        Watch the write collection and swap when full
+        """
+        counter_doc = self.db[self.collectionc].find_one({"_id": "uniqueFieldCounter"})
+        if counter_doc["sequence_value"] >= self.swap_cap:    # watch
+            return self.swap(generated)    # swap
+        return generated

wirehead-0.1.0/wirehead.egg-info/PKG-INFO

@@ -0,0 +1,4 @@
+Metadata-Version: 2.1
+Name: wirehead
+Version: 0.1.0
+License-File: LICENSE

wirehead-0.1.0/wirehead.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+LICENSE
+README.md
+setup.py
+wirehead/__init__.py
+wirehead/dataset.py
+wirehead/generator.py
+wirehead/manager.py
+wirehead.egg-info/PKG-INFO
+wirehead.egg-info/SOURCES.txt
+wirehead.egg-info/dependency_links.txt
+wirehead.egg-info/requires.txt
+wirehead.egg-info/top_level.txt

wirehead-0.1.0/wirehead.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

wirehead-0.1.0/wirehead.egg-info/requires.txt

@@ -0,0 +1,4 @@
+pymongo
+torch
+numpy
+PyYaml

wirehead-0.1.0/wirehead.egg-info/top_level.txt

@@ -0,0 +1 @@
+wirehead