Support rclone for sdxl

ctuning · Feb 2, 2024 · c4085bf · c4085bf
1 parent f9146ea
commit c4085bf
Show file tree

Hide file tree

Showing 10 changed files with 138 additions and 5 deletions.
diff --git a/cm-mlops/script/build-mlperf-inference-server-nvidia/_cm.yaml b/cm-mlops/script/build-mlperf-inference-server-nvidia/_cm.yaml
@@ -220,6 +220,8 @@ docker:
   base_image: nvcr.io/nvidia/mlperf/mlperf-inference:mlpinf-v3.1-cuda12.2-cudnn8.9-x86_64-ubuntu20.04-public
   docker_input_mapping:
     imagenet_path: IMAGENET_PATH
+    gptj_checkpoint_path: GPTJ_CHECKPOINT_PATH
+    criteo_preprocessed_path: CRITEO_PREPROCESSED_PATH
     results_dir: RESULTS_DIR
     submission_dir: SUBMISSION_DIR
     cudnn_tar_file_path: CM_CUDNN_TAR_FILE_PATH

diff --git a/cm-mlops/script/download-and-extract/_cm.json b/cm-mlops/script/download-and-extract/_cm.json
@@ -96,6 +96,14 @@
         "CM_DAE_EXTRACT_DOWNLOADED": "yes"
       }
     },
+    "rclone": {
+      "add_deps_recursive": {
+        "download-script": {
+          "tags": "_rclone"
+        }
+      },
+      "group": "download-tool"
+    },
     "gdown": {
       "add_deps_recursive": {
         "download-script": {

diff --git a/cm-mlops/script/download-file/_cm.json b/cm-mlops/script/download-file/_cm.json
@@ -63,6 +63,17 @@
       },
       "group": "download-tool"
     },
+    "rclone": {
+      "deps": [
+        {
+          "tags": "get,rclone"
+        }
+      ],
+      "env": {
+        "CM_DOWNLOAD_TOOL": "rclone"
+      },
+      "group": "download-tool"
+    },
     "url.#": {
       "env": {
         "CM_DOWNLOAD_URL": "#"

diff --git a/cm-mlops/script/download-file/customize.py b/cm-mlops/script/download-file/customize.py
@@ -65,6 +65,8 @@ def preprocess(i):
                 if j>0:
                     urltail=urltail[:j]
                 env['CM_DOWNLOAD_FILENAME'] = urltail
+            elif env.get('CM_DOWNLOAD_TOOL', '') == "rclone":
+                env['CM_DOWNLOAD_FILENAME'] = urltail
             else:
                 env['CM_DOWNLOAD_FILENAME'] = "index.html"
 
@@ -104,6 +106,11 @@ def preprocess(i):
         elif tool == "gdown":
             env['CM_DOWNLOAD_CMD'] = f"gdown {extra_download_options} {url}"
 
+        elif tool == "rclone":
+            if env.get('CM_RCLONE_CONFIG_CMD', '') != '':
+                env['CM_DOWNLOAD_CONFIG_CMD'] = env['CM_RCLONE_CONFIG_CMD']
+            env['CM_DOWNLOAD_CMD'] = f"rclone copy {url} ./{env['CM_DOWNLOAD_FILENAME']} -P"
+
         filename = env['CM_DOWNLOAD_FILENAME']
         env['CM_DOWNLOAD_DOWNLOADED_FILENAME'] = filename
 

diff --git a/cm-mlops/script/download-file/run.sh b/cm-mlops/script/download-file/run.sh
@@ -1,5 +1,11 @@
 #!/bin/bash
 
+if [[ -n ${CM_DOWNLOAD_CONFIG_CMD} ]]; then
+  echo ""
+  echo "${CM_DOWNLOAD_CONFIG_CMD}"
+  eval "${CM_DOWNLOAD_CONFIG_CMD}"
+fi
+
 if [ -e ${CM_DOWNLOAD_DOWNLOADED_PATH} ]; then
   if [[ "${CM_DOWNLOAD_CHECKSUM_CMD}" != "" ]]; then
     echo ""

diff --git a/cm-mlops/script/get-ml-model-gptj/_cm.json b/cm-mlops/script/get-ml-model-gptj/_cm.json
@@ -22,7 +22,7 @@
         "CM_EXTRACT_FINAL_ENV_NAME": "GPTJ_CHECKPOINT_PATH",
         "CM_EXTRACT_TO_FOLDER": "gpt-j"
       },
-      "tags": "download-and-extract,_wget",
+      "tags": "download-and-extract",
       "update_tags_from_env_with_prefix": {
         "_url.": [
           "CM_PACKAGE_URL"
@@ -34,7 +34,8 @@
       "force_cache": true,
       "names": [
         "dae"
-      ]
+      ],
+      "extra_cache_tags": "gptj,model"
     }
   ],
   "tags": [
@@ -174,6 +175,15 @@
         "CM_GPTJ_INTEL_MODEL": "yes"
       }
     },
+    "wget": {
+      "group": "download-tool",
+      "default": true,
+      "add_deps_recursive": {
+        "dae": {
+          "tags": "_wget"
+        }
+      }
+    },
     "pytorch,int8,intel": {
     },
     "pytorch,int4,intel": {

diff --git a/cm-mlops/script/get-ml-model-huggingface-zoo/_cm.json b/cm-mlops/script/get-ml-model-huggingface-zoo/_cm.json
@@ -17,6 +17,7 @@
     }
   ],
   "input_mapping": {
+    "download_path": "CM_DOWNLOAD_PATH",
     "model_filename": "CM_MODEL_ZOO_FILENAME",
     "env_key": "CM_MODEL_ZOO_ENV_KEY"
   },

diff --git a/cm-mlops/script/get-ml-model-huggingface-zoo/customize.py b/cm-mlops/script/get-ml-model-huggingface-zoo/customize.py
@@ -13,7 +13,9 @@ def preprocess(i):
 
     script_path = i['run_script_input']['path']
 
-    path = os.getcwd()
+    path = env.get('CM_DOWNLOAD_PATH', '')
+    if path == '':
+        path = os.getcwd()
 
     if env.get('CM_GIT_CLONE_REPO', '') != 'yes':
         run_cmd = env.get('CM_PYTHON_BIN_WITH_PATH') + " " +  os.path.join(script_path, 'download_model.py')

diff --git a/cm-mlops/script/get-ml-model-stable-diffusion/_cm.json b/cm-mlops/script/get-ml-model-stable-diffusion/_cm.json
@@ -10,7 +10,8 @@
     "CM_ML_MODEL_WEIGHT_TRANSFORMATIONS": "no"
   },
   "input_mapping": {
-    "checkpoint": "SDXL_CHECKPOINT_PATH"
+    "checkpoint": "SDXL_CHECKPOINT_PATH",
+    "download_path": "CM_DOWNLOAD_PATH"
   },
   "new_env_keys": [
     "CM_ML_MODEL_*",
@@ -21,6 +22,9 @@
       "enable_if_env": {
         "CM_TMP_REQUIRE_DOWNLOAD": [
           "yes"
+        ],
+        "CM_DOWNLOAD_TOOL": [
+          "git"
         ]
       },
       "env": {
@@ -34,6 +38,29 @@
         "hf-zoo"
       ],
       "tags": "get,ml-model,huggingface,zoo,_clone-repo,_model-stub.stabilityai/stable-diffusion-xl-base-1.0"
+    },
+    {
+      "env": {
+        "CM_EXTRACT_FINAL_ENV_NAME": "CM_ML_MODEL_PATH",
+        "CM_EXTRACT_TO_FOLDER": "stable-diffusion-xl-base-1.0"
+      },
+      "tags": "download-and-extract",
+      "update_tags_from_env_with_prefix": {
+        "_url.": [
+          "CM_DOWNLOAD_URL"
+        ]
+      },
+      "enable_if_env": {
+        "CM_TMP_REQUIRE_DOWNLOAD": [ "yes" ],
+        "CM_DOWNLOAD_TOOL": [
+          "rclone"
+        ]
+      },
+      "force_cache": true,
+      "extra_cache_tags": "stable-diffusion,sdxl,model",
+      "names": [
+        "dae"
+      ]
     }
   ],
   "tags": [
@@ -60,6 +87,14 @@
       },
       "group": "precision"
     },
+    "fp16": {
+      "env": {
+        "CM_ML_MODEL_INPUT_DATA_TYPES": "fp16",
+        "CM_ML_MODEL_PRECISION": "fp16",
+        "CM_ML_MODEL_WEIGHT_DATA_TYPES": "fp16"
+      },
+      "group": "precision"
+    },
     "int8": {
       "env": {
         "CM_ML_MODEL_INPUT_DATA_TYPES": "int8",
@@ -87,6 +122,58 @@
         "CM_ML_MODEL_WEIGHT_DATA_TYPES": "uint8"
       },
       "group": "precision"
+    },
+    "huggingface": {
+      "group": "download-source",
+      "default_variations": {
+        "download-tool": "git"
+      }
+    },
+    "mlcommons": {
+      "group": "download-source",
+      "default": true,
+      "default_variations": {
+        "download-tool": "rclone"
+      }
+    },
+    "git": {
+      "group": "download-tool",
+      "env": {
+        "CM_DOWNLOAD_TOOL": "git"
+      }
+    },
+    "wget": {
+      "group": "download-tool",
+      "env": {
+        "CM_DOWNLOAD_TOOL": "wget"
+      },
+      "adr": {
+        "dae": {
+          "tags": "_wget"
+        }
+      }
+    },
+    "rclone": {
+      "group": "download-tool",
+      "env": {
+        "CM_RCLONE_CONFIG": "rclone config create mlc-inference s3 provider=LyveCloud access_key_id=0LITLNQMHZALM5AK secret_access_key=YQKYTMBY23TMZHLOYFJKL5CHHS0CWYUC endpoint=s3.us-east-1.lyvecloud.seagate.com",
+        "CM_DOWNLOAD_TOOL": "rclone"
+      },
+      "adr": {
+        "dae": {
+          "tags": "_rclone"
+        }
+      }
+    },
+    "rclone,fp32": {
+      "env": {
+        "CM_DOWNLOAD_URL": "mlc-inference:mlcommons-inference-wg-s3/stable_diffusion_fp32"
+      }
+    },
+    "rclone,fp16": {
+      "env": {
+        "CM_DOWNLOAD_URL": "mlc-inference:mlcommons-inference-wg-s3/stable_diffusion_fp16"
+      }
     }
   }
 }
diff --git a/cm-mlops/script/get-ml-model-stable-diffusion/customize.py b/cm-mlops/script/get-ml-model-stable-diffusion/customize.py
@@ -18,7 +18,6 @@ def postprocess(i):
     env = i['env']
 
     env['SDXL_CHECKPOINT_PATH'] = env['CM_ML_MODEL_PATH']
-    env['CM_ML_MODEL_PATH'] = env['SDXL_CHECKPOINT_PATH']
     env['CM_GET_DEPENDENT_CACHED_PATH'] = env['CM_ML_MODEL_PATH']
 
     return {'return':0}