fixes

code/dataset_wrapper.py

@@ -0,0 +1,156 @@
+import sys
+import argparse
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
+import logging
+
+import lmdb
+import dotenv
+from tqdm import tqdm
+
+dotenv.load_dotenv()
+
+from vsm_datascience_common.cycle_database_connection.cycle_data import get_cycle_by_id
+from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
+
+from utils.dataset_creation import get_features, train_scalers, save_scalers, scale_item
+from utils.lmdb_utils import save_to_lmdb, load_from_lmdb
+from utils.utils import get_variable_from_module
+
+MAX_LMDB_SIZE_IN_MB = 200_000
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def process_id_wrapper(cycle_id: str, feature_config: dict, env: lmdb.Environment):
+    try:
+        cycle = get_cycle_by_id(cycle_id)
+        features = get_features(cycle, feature_config)
+        save_to_lmdb(env, key=str(cycle_id), dataset=features)
+    except:
+        pass
+
+
+def scaling_wrapper(key_batch: str,
+                    scalers: dict,
+                    max_lmdb_size_in_mb: int,
+                    lmdb_env_dir: str):
+    env = lmdb.open(lmdb_env_dir, readonly=True, lock=False)
+    scaled_items = list()
+    for key in key_batch:
+        item = load_from_lmdb(env, key)
+        scaled_item = scale_item(item, scalers)
+        scaled_items.append(scaled_item)
+    env.close()
+    env = lmdb.open(lmdb_env_dir, map_size=max_lmdb_size_in_mb * 1024 * 1024)
+    for key, scaled_item in zip(key_batch, scaled_items):
+        save_to_lmdb(env, key=key, dataset=scaled_item)
+    env.close()
+
+
+def create_dataset(model_configuration: dict,
+                   lmdb_root_dir: str,
+                   max_lmdb_size_in_mb: int,
+                   max_workers: int) -> None:
+    feature_config = model_configuration["feature_config"]
+
+    logger.info(f"Creating dataset for feature set {feature_config['feature_set_name']}")
+
+    # fetch valid cycle ids from database
+    valid_cycle_ids = [x["_id"] for x in get_cycles_collection().aggregate(
+        feature_config["filter_criteria_pipeline"] + [
+            {
+                "$project": {
+                    "_id": 1
+                }
+            }
+        ]
+    )]
+
+    logger.info(f"Fetched {len(valid_cycle_ids)} valid cycle ids from database")
+
+    env = lmdb.open(f"{lmdb_root_dir}/{feature_config['feature_set_name']}", map_size=max_lmdb_size_in_mb * 1024 * 1024)
+
+    # create features for items
+    logger.info(f"Creating features for {len(valid_cycle_ids)} cycles")
+    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+        list(tqdm(executor.map(partial(process_id_wrapper, feature_config=feature_config, env=env), valid_cycle_ids),
+                  total=len(valid_cycle_ids)))
+
+    # convert ids (here bson objectids) to keys for use in lmdb
+    keys = [str(cycle_id) for cycle_id in valid_cycle_ids]
+
+    # train scalers for featues
+    logger.info("Training scalers for features")
+    all_scalers = dict()
+    sample = load_from_lmdb(env, keys[0])
+    for feature_set in tqdm(feature_config["feature_sets"]):
+        for feature in feature_config[feature_set]:
+            print(f"Training scalers for feature {feature['name']} in feature set {feature_set}")
+            scaler = train_scalers(feature_set, feature["name"], feature["scaler"], sample, env)
+            if feature_set not in all_scalers:
+                all_scalers[feature_set] = dict()
+            all_scalers[feature_set] = all_scalers[feature_set] | scaler
+
+    # save scalers
+    logger.info("Saving scalers to disk")
+    scaler_dir = f"{lmdb_root_dir}/{feature_config['feature_set_name']}/scalers"
+    save_scalers(all_scalers, scaler_dir)
+
+    # creat scaling batches for less burden on lmdb
+    batch_size = 1_000
+    key_batches = [keys[i:i + batch_size] for i in range(0, len(keys), batch_size)]
+
+    # scale items
+    logger.info("Scaling items")
+    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+        list(tqdm(executor.map(partial(scaling_wrapper, scalers=all_scalers,
+                                       max_lmdb_size_in_mb=max_lmdb_size_in_mb,
+                                       lmdb_env_dir=f"{lmdb_root_dir}/{feature_config['feature_set_name']}"),
+                               key_batches),
+                  total=len(key_batches)))
+
+    logger.info(f"Dataset creation finished. LMDB saved in {lmdb_root_dir}/{feature_config['feature_set_name']}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Dataset Creation Wrapper")
+    parser.add_argument("--model_config_module",
+                        type=str,
+                        required=True,
+                        help="Path to the model config module")
+    parser.add_argument("--model_config_variable",
+                        type=str,
+                        required=False,
+                        default="model_configuration",
+                        help="Name of the model config variable, default is 'model_configuration'")
+    parser.add_argument("--lmdb_dir",
+                        type=str,
+                        required=False,
+                        default="./lmdb_datasets",
+                        help="Path to the lmdb directory")
+    parser.add_argument("--lmdb_size",
+                        type=int,
+                        required=False,
+                        default=MAX_LMDB_SIZE_IN_MB,
+                        help="Size of the lmdb in MB, default is 200_000")
+    parser.add_argument("--max_workers",
+                        type=int,
+                        required=False,
+                        default=None,
+                        help="Number of workers for multiprocessing, default is None (use all available cores)")
+
+    args = parser.parse_args()
+
+    # import and load the model config
+    model_configuration = get_variable_from_module(
+        module_path=args.model_config_module,
+        variable_name=args.model_config_variable
+    )
+
+    create_dataset(model_configuration,
+                   lmdb_root_dir=args.lmdb_dir,
+                   max_lmdb_size_in_mb=args.lmdb_size,
+                   max_workers=args.max_workers)