[to #50129142]feat: support pre build torch extension build image, first extension megatron_util

.dev_scripts/build_image.sh

@@ -132,7 +132,7 @@ echo -e "Building image with:\npython$python_version\npytorch$torch_version\nten
 docker_file_content=`cat docker/Dockerfile.ubuntu`
 if [ "$is_ci_test" != "True" ]; then
     echo "Building ModelScope lib, will install ModelScope lib to image"
-    docker_file_content="${docker_file_content} \nRUN pip install --no-cache-dir  modelscope==$modelscope_version -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html"
+    docker_file_content="${docker_file_content} \nRUN pip install --no-cache-dir modelscope==$modelscope_version -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html"
 fi
 echo "$is_dsw"
 if [ "$is_dsw" == "False" ]; then
@@ -140,6 +140,8 @@ if [ "$is_dsw" == "False" ]; then
 else
     echo "Building dsw image will need set ModelScope lib cache location."
     docker_file_content="${docker_file_content} \nENV MODELSCOPE_CACHE=/mnt/workspace/.cache/modelscope"
+    # pre compile extension
+    docker_file_content="${docker_file_content} \nRUN python -c 'from modelscope.utils.pre_compile import pre_compile_all;pre_compile_all()'"
 fi
 if [ "$is_ci_test" == "True" ]; then
     echo "Building CI image, uninstall modelscope"

modelscope/pipelines/cv/card_detection_pipeline.py

@@ -6,8 +6,8 @@ from modelscope.models.base.base_model import Model
 from modelscope.pipelines.base import Pipeline
 from modelscope.pipelines.builder import PIPELINES
 from modelscope.utils.constant import Tasks
+from modelscope.utils.input_output_typing import Image
 from modelscope.utils.logger import get_logger
-from modelscope.utils.typing import Image
 
 logger = get_logger()
 

modelscope/pipelines/cv/face_detection_pipeline.py

@@ -16,8 +16,8 @@ from modelscope.pipelines.builder import PIPELINES
 from modelscope.preprocessors import LoadImage
 from modelscope.utils.config import Config
 from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.input_output_typing import Image
 from modelscope.utils.logger import get_logger
-from modelscope.utils.typing import Image
 
 logger = get_logger()
 

modelscope/utils/pre_compile.py

@@ -0,0 +1,27 @@
+import os
+
+import torch
+
+from modelscope.utils.megatron_utils import init_megatron_util
+
+
+def pre_compile_megatron_util():
+    dummy_megatron_cfg = {
+        'tensor_model_parallel_size': 1,
+        'world_size': 1,
+        'distributed_backend': 'nccl',
+        'seed': 42,
+    }
+    os.environ['MASTER_PORT'] = '39501'
+    init_megatron_util(dummy_megatron_cfg)
+
+
+def pre_compile_all():
+    if torch.cuda.is_available():  # extension require cuda.
+        pre_compile_megatron_util()
+
+    # extension for all platform.
+
+
+if __name__ == '__main__':
+    pre_compile_all()