Repository: mlcommons/GaNDLF

Issue:
The inference loop is currently using `modelDir` as the location to write outputs, which can be problematic for deployment
**Describe the bug**
When deploying a model (regardless of whether that is through a container or any other means), it is prudent to assume that the process will not have write access to `modelDir`. 

**To Reproduce**
N.A.

**Expected behavior**
Split the `outputDir_or_optimizedModel` variable in the [inference_loop](https://github.com/mlcommons/GaNDLF/blob/master/GANDLF/compute/inference_loop.py#L38) into `optimizedModel` and `outputDir`. We _could_ default `outputDir` to the location of `optimizedModel` if not present.

**Screenshots**
N.A.

**GaNDLF Version**
<!-- Put the output of the following command:
python -c 'import GANDLF as g;print(g.__version__)'
-->
0.0.16-dev

**Desktop (please complete the following information):**
N.A.

**Additional context**
N.A.


Discussion:

Resolving patch:
diff --git a/GANDLF/cli/main_run.py b/GANDLF/cli/main_run.py
--- a/GANDLF/cli/main_run.py
+++ b/GANDLF/cli/main_run.py
@@ -1,4 +1,4 @@
-import os, pickle, sys
+import os, pickle
 from pathlib import Path
 
 from GANDLF.training_manager import TrainingManager, TrainingManager_split
@@ -7,21 +7,24 @@
 from GANDLF.utils import populate_header_in_parameters, parseTrainingCSV
 
 
-def main_run(data_csv, config_file, output_dir, train_mode, device, resume, reset):
+def main_run(
+    data_csv, config_file, model_dir, train_mode, device, resume, reset, output_dir=None
+):
     """
     Main function that runs the training and inference.
 
     Args:
         data_csv (str): The CSV file of the training data.
         config_file (str): The YAML file of the training configuration.
-        output_dir (str): The output directory.
+        model_dir (str): The model directory; for training, model is written out here, and for inference, trained model is expected here.
         train_mode (bool): Whether to train or infer.
         device (str): The device type.
         resume (bool): Whether the previous run will be resumed or not.
         reset (bool): Whether the previous run will be reset or not.
+        output_dir (str): The output directory for the inference session.
 
-    Raises:
-        ValueError: Parameter check from previous run.
+    Returns:
+        None
     """
     file_data_full = data_csv
     model_parameters = config_file
@@ -30,7 +33,7 @@ def main_run(data_csv, config_file, output_dir, train_mode, device, resume, rese
     parameters["device_id"] = -1
     # in case the data being passed is already processed, check if the previous parameters exists,
     # and if it does, compare the two and if they are the same, ensure no preprocess is done.
-    model_parameters_prev = os.path.join(os.path.dirname(output_dir), "parameters.pkl")
+    model_parameters_prev = os.path.join(os.path.dirname(model_dir), "parameters.pkl")
     if train_mode:
         if not (reset) or not (resume):
             print(
@@ -39,19 +42,15 @@ def main_run(data_csv, config_file, output_dir, train_mode, device, resume, rese
             )
             if os.path.exists(model_parameters_prev):
                 parameters_prev = pickle.load(open(model_parameters_prev, "rb"))
-                if parameters != parameters_prev:
-                    raise ValueError(
-                        "The parameters are not the same as the ones stored in the previous run, please re-check."
-                    )
-
-    parameters["output_dir"] = output_dir
+                assert (
+                    parameters == parameters_prev
+                ), "The parameters are not the same as the ones stored in the previous run, please re-check."
+        parameters["output_dir"] = model_dir
+        Path(parameters["output_dir"]).mkdir(parents=True, exist_ok=True)
 
     if "-1" in device:
         device = "cpu"
 
-    if train_mode:
-        Path(parameters["output_dir"]).mkdir(parents=True, exist_ok=True)
-
     # parse training CSV
     if "," in file_data_full:
         # training and validation pre-split
@@ -104,7 +103,8 @@ def main_run(data_csv, config_file, output_dir, train_mode, device, resume, rese
         else:
             InferenceManager(
                 dataframe=data_full,
-                outputDir=parameters["output_dir"],
+                modelDir=model_dir,
+                outputDir=output_dir,
                 parameters=parameters,
                 device=device,
             )
diff --git a/GANDLF/compute/inference_loop.py b/GANDLF/compute/inference_loop.py
--- a/GANDLF/compute/inference_loop.py
+++ b/GANDLF/compute/inference_loop.py
@@ -36,7 +36,7 @@ def applyCustomColorMap(im_gray):
 
 
 def inference_loop(
-    inferenceDataFromPickle, device, parameters, outputDir_or_optimizedModel
+    inferenceDataFromPickle, device, parameters, modelDir, outputDir=None
 ):
     """
     The main training loop.
@@ -45,7 +45,8 @@ def inference_loop(
         inferenceDataFromPickle (pandas.DataFrame): The data to use for inference.
         device (str): The device to perform computations on.
         parameters (dict): The parameters dictionary.
-        outputDir_or_optimizedModel (str): The output directory or optimized model file.
+        modelDir (str): The path to the directory containing the model to be used for inference.
+        outputDir (str): The path to the directory where the output of the inference session will be stored.
     """
     # Defining our model here according to parameters mentioned in the configuration file
     print("Current model type : ", parameters["model"]["type"])
@@ -71,14 +72,14 @@ def inference_loop(
     main_dict = None
     if parameters["model"]["type"] == "torch":
         # Loading the weights into the model
-        if os.path.isdir(outputDir_or_optimizedModel):
+        if os.path.isdir(modelDir):
             files_to_check = [
                 os.path.join(
-                    outputDir_or_optimizedModel,
+                    modelDir,
                     str(parameters["model"]["architecture"]) + best_model_path_end,
                 ),
                 os.path.join(
-                    outputDir_or_optimizedModel,
+                    modelDir,
                     str(parameters["model"]["architecture"]) + latest_model_path_end,
                 ),
             ]
@@ -96,13 +97,13 @@ def inference_loop(
         model.eval()
     elif parameters["model"]["type"].lower() == "openvino":
         # Loading the executable OpenVINO model
-        if os.path.isdir(outputDir_or_optimizedModel):
+        if os.path.isdir(modelDir):
             xml_to_check = os.path.join(
-                outputDir_or_optimizedModel,
+                modelDir,
                 str(parameters["model"]["architecture"]) + "_best.xml",
             )
             bin_to_check = os.path.join(
-                outputDir_or_optimizedModel,
+                modelDir,
                 str(parameters["model"]["architecture"]) + "_best.bin",
             )
             if not os.path.isfile(xml_to_check):
@@ -177,9 +178,7 @@ def inference_loop(
             level_width, level_height = os_image.level_dimensions[
                 parameters["slide_level"]
             ]
-            subject_dest_dir = os.path.join(
-                outputDir_or_optimizedModel, str(subject_name)
-            )
+            subject_dest_dir = os.path.join(outputDir, str(subject_name))
             Path(subject_dest_dir).mkdir(parents=True, exist_ok=True)
 
             try:
diff --git a/GANDLF/inference_manager.py b/GANDLF/inference_manager.py
--- a/GANDLF/inference_manager.py
+++ b/GANDLF/inference_manager.py
@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 import pandas as pd
 import torch
 import torch.nn.functional as F
@@ -7,13 +8,32 @@
 from GANDLF.utils import get_unique_timestamp
 
 
-def InferenceManager(dataframe, outputDir, parameters, device):
+def InferenceManager(dataframe, modelDir, parameters, device, outputDir=None):
     """
-    This function takes in a dataframe, with some other parameters and performs the inference
+    This function takes in a dataframe, with some other parameters and performs the inference on the data in the dataframe.
+
+    Args:
+        dataframe (pandas.DataFrame): The dataframe containing the data to be used for inference.
+        modelDir (str): The path to the directory containing the model to be used for inference.
+        outputDir (str): The path to the directory where the output of the inference will be stored.
+        parameters (dict): The dictionary containing the parameters for the inference.
+        device (str): The device type.
+
+    Returns:
+        None
     """
     # get the indeces for kfold splitting
     inferenceData_full = dataframe
 
+    # if outputDir is not provided, create a new directory with a unique timestamp
+    if outputDir is None:
+        outputDir = os.path.join(modelDir, get_unique_timestamp())
+        print(
+            "Output directory not provided, creating a new directory with a unique timestamp: ",
+            outputDir,
+        )
+    Path(outputDir).mkdir(parents=True, exist_ok=True)
+
     # # initialize parameters for inference
     if not ("weights" in parameters):
         parameters["weights"] = None  # no need for loss weights for inference
@@ -22,24 +42,24 @@ def InferenceManager(dataframe, outputDir, parameters, device):
 
     n_folds = parameters["nested_training"]["validation"]
 
-    outputDir_split = [outputDir]
-    if "," in outputDir:
-        outputDir_split = outputDir.split(",")
+    modelDir_split = [modelDir]
+    if "," in modelDir:
+        modelDir_split = modelDir.split(",")
 
     averaged_probs_list = []
-    for current_outputDir in outputDir_split:
+    for current_modelDir in modelDir_split:
         fold_dirs = []
         if n_folds > 1:
-            directories = sorted(os.listdir(current_outputDir))
+            directories = sorted(os.listdir(current_modelDir))
             for d in directories:
                 if d.isdigit():
-                    fold_dirs.append(os.path.join(current_outputDir, d, ""))
+                    fold_dirs.append(os.path.join(current_modelDir, d, ""))
         else:
-            fold_dirs = [current_outputDir]
+            fold_dirs = [current_modelDir]
 
         # this is for the case where inference is happening using a single model
         if len(fold_dirs) == 0:
-            fold_dirs = [current_outputDir]
+            fold_dirs = [current_modelDir]
 
         probs_list = []
         class_list = None
@@ -53,9 +73,10 @@ def InferenceManager(dataframe, outputDir, parameters, device):
             parameters["current_fold_dir"] = fold_dir
             inference_loop(
                 inferenceDataFromPickle=inferenceData_full,
-                outputDir_or_optimizedModel=fold_dir,
+                modelDir=fold_dir,
                 device=device,
                 parameters=parameters,
+                outputDir=outputDir,
             )
 
             if is_classification:
@@ -88,12 +109,10 @@ def InferenceManager(dataframe, outputDir, parameters, device):
         averaged_probs_df.PredictedClass = [
             class_list[a] for a in averaged_probs_across_models.argmax(1)
         ]
-        filepath_to_save = os.path.join(
-            outputDir_split[0], "final_preds_and_avg_probs.csv"
-        )
+        filepath_to_save = os.path.join(outputDir, "final_preds_and_avg_probs.csv")
         if os.path.isfile(filepath_to_save):
             filepath_to_save = os.path.join(
-                outputDir_split[0],
+                outputDir,
                 "final_preds_and_avg_probs" + get_unique_timestamp() + ".csv",
             )
         averaged_probs_df.to_csv(filepath_to_save, index=False)


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -651,10 +651,18 @@ def test_train_resume_inference_classification_rad_3d(device):
     parameters["output_dir"] = outputDir  # this is in inference mode
     InferenceManager(
         dataframe=training_data,
-        outputDir=outputDir,
+        modelDir=outputDir,
         parameters=parameters,
         device=device,
     )
+    # test the case where outputDir is explicitly provided to InferenceManager
+    InferenceManager(
+        dataframe=training_data,
+        modelDir=outputDir,
+        parameters=parameters,
+        device=device,
+        outputDir=os.path.join(outputDir, get_unique_timestamp()),
+    )
 
     print("passed")
 
@@ -693,7 +701,7 @@ def test_train_inference_optimize_classification_rad_3d(device):
         parameters["output_dir"] = outputDir  # this is in inference mode
         InferenceManager(
             dataframe=training_data,
-            outputDir=outputDir,
+            modelDir=outputDir,
             parameters=parameters,
             device=device,
         )
@@ -738,7 +746,7 @@ def test_train_inference_optimize_segmentation_rad_2d(device):
         parameters["output_dir"] = outputDir  # this is in inference mode
         InferenceManager(
             dataframe=training_data,
-            outputDir=outputDir,
+            modelDir=outputDir,
             parameters=parameters,
             device=device,
         )
@@ -801,7 +809,7 @@ def test_train_inference_classification_with_logits_single_fold_rad_3d(device):
     parameters["model"]["onnx_export"] = False
     InferenceManager(
         dataframe=training_data,
-        outputDir=outputDir,
+        modelDir=outputDir,
         parameters=parameters,
         device=device,
     )
@@ -844,7 +852,7 @@ def test_train_inference_classification_with_logits_multiple_folds_rad_3d(device
     parameters["output_dir"] = outputDir  # this is in inference mode
     InferenceManager(
         dataframe=training_data,
-        outputDir=outputDir + "," + outputDir,
+        modelDir=outputDir + "," + outputDir,
         parameters=parameters,
         device=device,
     )
@@ -1949,7 +1957,7 @@ def test_train_inference_segmentation_histology_2d(device):
     inference_data.drop(index=inference_data.index[-1], axis=0, inplace=True)
     InferenceManager(
         dataframe=inference_data,
-        outputDir=modelDir,
+        modelDir=modelDir,
         parameters=parameters,
         device=device,
     )
@@ -2082,17 +2090,30 @@ def test_train_inference_classification_histology_large_2d(device):
         parameters["model"]["type"] = model_type
         InferenceManager(
             dataframe=inference_data,
-            outputDir=modelDir,
+            modelDir=modelDir,
             parameters=parameters,
             device=device,
         )
-        # if 'predictions.csv' are not found, give error
-        output_subject_dir = os.path.join(modelDir, str(input_df["SubjectID"][0]))
-        assert (
-            os.path.exists(os.path.join(output_subject_dir, "predictions.csv")) is True
-        )
-        # ensure previous results are removed
-        shutil.rmtree(output_subject_dir)
+        all_folders_in_modelDir = os.listdir(modelDir)
+        for folder in all_folders_in_modelDir:
+            output_subject_dir = os.path.join(modelDir, folder)
+            if os.path.isdir(output_subject_dir):
+                # check in the default outputDir that's created - this is based on a unique timestamp
+                if folder != "output_validation":
+                    # if 'predictions.csv' are not found, give error
+                    assert (
+                        os.path.exists(
+                            os.path.join(
+                                output_subject_dir,
+                                str(input_df["SubjectID"][0]),
+                                "predictions.csv",
+                            )
+                        )
+                        is True,
+                        "predictions.csv not found",
+                    )
+    # ensure previous results are removed
+    sanitize_outputDir()
 
     for file in files_to_delete:
         os.remove(file)
@@ -2178,7 +2199,7 @@ def test_train_inference_classification_histology_2d(device):
         parameters["model"]["type"] = model_type
         InferenceManager(
             dataframe=inference_data,
-            outputDir=modelDir,
+            modelDir=modelDir,
             parameters=parameters,
             device=device,
         )