Refactor build_wheel to support oneflowinc images (#5427)

* refactor_build_wheel_to_support_oneflowinc_images

* address review
Co-authored-by: Noneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

.github/actions/whl/action.yml

@@ -31,6 +31,5 @@ runs:
             --wheel_house_dir=${tmp_dir}/wheelhouse \
             --oneflow_src_dir=${src_dir} ${{ inputs.extra_flags }} \
             --retry=1 \
-            --skip_img \
             --extra_docker_args "${extra_docker_args}"
       shell: bash

README.md

@@ -93,10 +93,18 @@
   Please refer to [this repo](https://github.com/Oneflow-Inc/conda-env)
 
 - #### Option 2: Build in docker container (recommended)
+  - Pull a docker image:
+
+    ```
+    docker pull oneflowinc/oneflow-manylinux2014-cuda10.2:0.1
+    ```
+
+    All images available : https://hub.docker.com/u/oneflowinc
+
   - In the root directory of OneFlow source code, run:
 
     ```
-    python3 docker/package/manylinux/build_wheel.py
+    python3 docker/package/manylinux/build_wheel.py --python_version=3.6
     ```
 
     This should produce `.whl` files in the directory `wheelhouse`

docker/package/manylinux/build_wheel.py

@@ -46,6 +46,7 @@ def build_img(
     use_tuna,
     use_system_proxy,
     img_tag,
+    dry,
 ):
     cudnn_version = 7
     if str(cuda_version).startswith("11"):
@@ -67,7 +68,8 @@ def build_img(
     proxy_build_arg = get_proxy_build_args() if use_system_proxy else ""
     cmd = f"docker build -f docker/package/manylinux/Dockerfile {proxy_build_arg} {tuna_build_arg} --build-arg from={from_img} -t {img_tag} ."
     print(cmd)
-    subprocess.check_call(cmd, cwd=oneflow_src_dir, shell=True)
+    if dry == False:
+        subprocess.check_call(cmd, cwd=oneflow_src_dir, shell=True)
 
 
 def common_cmake_args(cache_dir=None, extra_oneflow_cmake_args=None):
@@ -310,9 +312,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--wheel_house_dir", type=str, required=False, default=default_wheel_house_dir,
     )
-    parser.add_argument(
-        "--python_version", type=str, required=False, default="3.6, 3.7, 3.8",
-    )
+    parser.add_argument("--python_version", type=str, required=True)
     parser.add_argument(
         "--cuda_version", type=str, required=False, default="10.2",
     )
@@ -340,6 +340,9 @@ if __name__ == "__main__":
     parser.add_argument(
         "--skip_audit", default=False, action="store_true", required=False
     )
+    parser.add_argument(
+        "--build_img", default=False, action="store_true", required=False
+    )
     parser.add_argument(
         "--use_tuna", default=False, action="store_true", required=False
     )
@@ -357,6 +360,8 @@ if __name__ == "__main__":
     parser.add_argument("--bash", default=False, action="store_true", required=False)
     parser.add_argument("--retry", default=0, type=int)
     args = parser.parse_args()
+    if args.skip_img:
+        "Arg skip_img is deprecated. Setting it has no effect. If you want to build image, use --build_img"
     print("args.extra_oneflow_cmake_args", args.extra_oneflow_cmake_args)
     assert args.package_name
     extra_oneflow_cmake_args = " ".join(
@@ -383,7 +388,6 @@ if __name__ == "__main__":
 
         def build():
             img_tag = None
-            skip_img = args.skip_img
             img_prefix = f"oneflow-manylinux2014-cuda{cuda_version}"
             user = getpass.getuser()
             versioned_img_tag = f"{img_prefix}:0.1"
@@ -401,25 +405,30 @@ if __name__ == "__main__":
                 enforced_oneflow_cmake_args += (
                     ' -DBAZEL_ENV_ARGS="BAZEL_LINKLIBS=-l%:libstdc++.a"'
                 )
-            user_img_tag = f"{img_prefix}:{user}"
             extra_docker_args = args.extra_docker_args
             if "--name" not in extra_docker_args:
                 extra_docker_args += (
                     f" --name run-by-{getpass.getuser()}-{str(uuid.uuid4())}"
                 )
-            if args.custom_img_tag:
+            user_img_tag = f"{img_prefix}:{user}"
+            inc_img_tag = f"oneflowinc/{versioned_img_tag}"
+            img_tag = inc_img_tag
+            if args.build_img:
+                img_tag = user_img_tag
+            elif args.custom_img_tag:
                 img_tag = args.custom_img_tag
-                skip_img = True
-            elif skip_img:
-                assert is_img_existing(
-                    versioned_img_tag
-                ), f"img not found: {versioned_img_tag}"
-                img_tag = versioned_img_tag
             else:
-                img_tag = user_img_tag
+                if is_img_existing(versioned_img_tag):
+                    img_tag = versioned_img_tag
+                elif is_img_existing(inc_img_tag):
+                    img_tag = inc_img_tag
+                else:
+                    raise ValueError(
+                        f"img not found, please run 'docker pull {inc_img_tag}'"
+                    )
             assert img_tag is not None
             print("using", img_tag)
-            if skip_img == False:
+            if args.build_img:
                 build_img(
                     cuda_version,
                     args.oneflow_src_dir,
@@ -427,6 +436,7 @@ if __name__ == "__main__":
                     args.use_tuna,
                     args.use_system_proxy,
                     img_tag,
+                    args.dry,
                 )
             bash_args = ""
             if args.xla:
@@ -461,6 +471,8 @@ gcc --version
                     assert len(cuda_versions) == 1
                     sub_dir += "-cpu"
                 cache_dir = os.path.join(cache_dir, sub_dir)
+            if args.build_img:
+                return
             if args.skip_third_party == False:
                 build_third_party(
                     img_tag,