Create export path if not present (#3331)

Co-authored-by: Ankith Gunapal <[email protected]>
pytorch · Oct 3, 2024 · e8879c1 · e8879c1
1 parent 9d10087
commit e8879c1
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/model-archiver/model_archiver/model_packaging_utils.py b/model-archiver/model_archiver/model_packaging_utils.py
@@ -43,6 +43,8 @@
 MANIFEST_FILE_NAME = "MANIFEST.json"
 MAR_INF = "MAR-INF"
 
+logger = logging.getLogger(__file__)
+
 
 class ModelExportUtils(object):
     """
@@ -366,7 +368,12 @@ def check_model_name_regex_or_exit(model_name):
     @staticmethod
     def validate_inputs(model_name, export_path):
         ModelExportUtils.check_model_name_regex_or_exit(model_name)
-        if not os.path.isdir(os.path.abspath(export_path)):
+        if not os.path.exists(os.path.abspath(export_path)):
+            logger.warning(
+                f"Export directory ({export_path}) does not exist. Creating..."
+            )
+            os.makedirs(os.path.abspath(export_path))
+        elif not os.path.isdir(os.path.abspath(export_path)):
             raise ModelArchiverError(
                 "Given export-path {} is not a directory. "
                 "Point to a valid export-path directory.".format(export_path)