added code

code/utils/training_utils.py

@@ -1,3 +1,4 @@
+import math
 import os
 import pickle
 import random
@@ -14,7 +15,7 @@ from tqdm import tqdm
 
 from utils.lmdb_utils import get_lmdb_keys
 from utils.utils import get_config_id
-from utils.data_utils import produce_window_batches
+from utils.data_utils import produce_window_batches, get_collated_batch_for_key
 
 from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
 
@@ -26,11 +27,11 @@ def get_training_config(base_config: dict, model_config: dict):
     except KeyError:
         raise Exception("Model class missing in training or model config")
 
-    # training_config_id = get_config_id(base_config)
+    training_config_id = get_config_id(base_config)
     # hash = training_config_id[-5:]
 
-    timestamp = datetime.now().strftime("%Y%m%d-%H%M")
-    training_config_id = f"{timestamp}"
+    # timestamp = datetime.now().strftime("%Y%m%d-%H%M")
+    # training_config_id = f"{timestamp}"
 
     training_dir = os.path.abspath(f"{model_config['model_dir']}/trainings/{training_config_id}")
 
@@ -49,6 +50,10 @@ def get_training_config(base_config: dict, model_config: dict):
         with open(f"{training_dir}/training_configuration.pickle", "rb") as f:
             training_config = pickle.load(f)
 
+        # update paths
+        training_config["training_dir"] = training_dir
+        training_config["model_dir"] = model_config["model_dir"]
+
     return training_config
 
 
@@ -68,6 +73,78 @@ def get_training_config_from_file(training_dir: str,
     return training_config
 
 
+def load_model_for_usage(model_configuration: dict,
+                         training_configuration: dict,
+                         sample_id: str) -> tuple[nn.Module, dict]:
+    """
+    Load the model associated with the given model configuration for the usage on the current node, given its training configuration
+
+
+    If the model does not exist, create it
+    Args:
+        model_configuration: configuration of the model
+        training_configuration: training configuration
+        sample_id: id of a sample data item
+    Returns:
+        model: model
+        training_config: training configuration with updated variables such as batch size
+
+    """
+
+    # get computation rank
+    if torch.distributed.is_initialized():
+        local_rank = torch.distributed.get_rank()
+        torch.cuda.set_device(local_rank)
+        device = torch.device(f"cuda:{local_rank}")
+    else:
+        local_rank = 0
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    dataset_dir = model_configuration["feature_config"]["dataset_dir"]
+    lmdb_env = lmdb.open(dataset_dir, readonly=True)
+
+    try:
+        load_fn = model_configuration["model_load_fn"]
+        model = load_fn(model_configuration=model_configuration,
+                        training_configuration=training_configuration,
+                        sample_key=sample_id,
+                        device=device,
+                        lmdb_env=lmdb_env)
+    except FileNotFoundError:
+        model_creation_fn = model_configuration["model_creation_fn"]
+        model = model_creation_fn(
+            model_configuration=model_configuration,
+            sample_key=sample_id,
+            lmdb_env=lmdb_env
+        )
+
+    # estimate batch size
+    sample_batch = get_collated_batch_for_key(sample_id, model_configuration, lmdb_env=lmdb_env)
+    if sample_batch is None or len(sample_batch) == 0:
+        raise ValueError(
+            f"Rank {local_rank}: Could not get sample input for key {sample_id} in dataset {dataset_dir}. "
+            f"Please check if the key exists in the LMDB dataset.")
+    sample_input = sample_batch[0][0]  # get the first item in the batch
+    input_shape = sample_input.shape
+    # send model to proper device
+    model.to(device)
+    batch_size = find_max_batch_size(model,
+                                     input_shape,
+                                     device=device)
+    # update max lr in training configuration, use square root scaling law based on initial batch size
+    training_configuration["learning_parameters"]["learning_rate"] = training_configuration["learning_parameters"][
+                                                                         "max_lr"] * math.sqrt(
+        batch_size / training_configuration["batch_size"])
+
+    # update batch size in training configuration
+    training_configuration["batch_size"] = batch_size
+
+    # close lmdb env
+    lmdb_env.close()
+
+    return model, training_configuration
+
+
 def weight_init(m):
     """
     Usage:
@@ -259,3 +336,44 @@ def get_data_ids(model_configuration: dict,
         test_ids = test_ids[:test_lim]
 
     return train_ids, val_ids, test_ids
+
+
+def find_max_batch_size(model,
+                        input_shape,
+                        device,
+                        dtype=torch.float32,
+                        max_batch=2048):
+    """
+    Find the maximum batch size that can be processed by the model without running out of memory.
+
+    Model must already be moved to the appropriate device (e.g., GPU).
+    Args:
+        model: model to test
+        input_shape: shape of the input tensor (excluding batch size)
+        device: device to work on
+        dtype: data type of the input tensor (default: torch.float32)
+        max_batch: maximum batch size to test (default: 2048)
+
+    Returns:
+        last_successful: the largest batch size that did not cause an out of memory error
+
+    """
+    batch_size = 1
+    last_successful = 1
+
+    while last_successful < max_batch:
+        try:
+            inputs = torch.randn((batch_size, *input_shape), device=device, dtype=dtype)
+            outputs = model(inputs)
+            loss = outputs.sum()
+            loss.backward()  # simulate training
+            torch.cuda.empty_cache()
+            last_successful = batch_size
+            batch_size *= 2
+        except RuntimeError as e:
+            if 'out of memory' in str(e):
+                torch.cuda.empty_cache()
+                break
+            else:
+                raise e
+    return last_successful