From bcc10ab73c8b922a525c675a52e09a5c429da7bf Mon Sep 17 00:00:00 2001
From: zhoufeng <zhoufeng54@huawei.com>
Date: Sun, 22 Nov 2020 15:50:30 +0800
Subject: [PATCH] check version of te/topi/hccl wheel package while "import
 mindspore"ing

Signed-off-by: zhoufeng <zhoufeng54@huawei.com>
---
 mindspore/_check_deps_version.py | 74 ++++++++++++++++++++++++++++++++
 mindspore/_check_version.py      | 17 +++++++-
 2 files changed, 90 insertions(+), 1 deletion(-)
 create mode 100644 mindspore/_check_deps_version.py

diff --git a/mindspore/_check_deps_version.py b/mindspore/_check_deps_version.py
new file mode 100644
index 0000000000..2ee4c88c62
--- /dev/null
+++ b/mindspore/_check_deps_version.py
@@ -0,0 +1,74 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""dependency package version check"""
+from argparse import ArgumentParser
+
+def parse_args():
+    """
+    parse args .
+
+    Args:
+
+    Returns:
+        args.
+
+    Examples:
+        >>> parse_args()
+    """
+    parser = ArgumentParser(description="MindSpore dependency packages version checker.")
+    parser.add_argument("--mindspore_version", type=str, help="MindSpore version.")
+    parser.add_argument("--supported_version", type=str, action='append', help="Supported environment version.")
+    args = parser.parse_args()
+    return args
+
+def check_deps_version(mindspore_version, supported_version):
+    """
+       check te/hccl/topi version
+
+    Args:
+        mindspore_version (str): this mindspore package version
+        supported_version (str list): supported Ascend 910 AI software package version by this mindspore package
+
+    Returns:
+        void
+    """
+    try:
+        import te
+        v = te.__version__
+        if v not in supported_version:
+            print(f"MindSpore version {mindspore_version} and \"te\" wheel package version {v} does not "
+                  "match, reference to the match info on: https://www.mindspore.cn/install")
+        import topi
+        v = topi.__version__
+        if v not in supported_version:
+            print(f"MindSpore version {mindspore_version} and \"topi\" wheel package version {v} does not "
+                  "match, reference to the match info on: https://www.mindspore.cn/install")
+        from hccl import sys_version as hccl_version
+        v = hccl_version.__sys_version__
+        if v not in supported_version:
+            print(f"MindSpore version {mindspore_version} and \"hccl\" wheel package version {v} does not "
+                  "match, reference to the match info on: https://www.mindspore.cn/install")
+    except ImportError as e:
+        print("CheckFailed: ", e.args)
+        print("Minspore relies on the 3 whl packages of \"te\", \"topi\" and \"hccl\" in the \"fwkacllib\" "
+              "folder of the Ascend 910 AI software package, please check whether they are installed "
+              "correctly or not, reference to the match info on: https://www.mindspore.cn/install")
+
+def main():
+    args = parse_args()
+    check_deps_version(args.mindspore_version, args.supported_version)
+
+if __name__ == "__main__":
+    main()
diff --git a/mindspore/_check_version.py b/mindspore/_check_version.py
index 3c3025a750..136cbd8afb 100644
--- a/mindspore/_check_version.py
+++ b/mindspore/_check_version.py
@@ -15,6 +15,7 @@
 """version and config check"""
 import os
 import sys
+import subprocess
 from pathlib import Path
 from abc import abstractmethod, ABCMeta
 from packaging import version
@@ -188,6 +189,17 @@ class AscendEnvChecker(EnvChecker):
             logger.warning(f"MindSpore version {__version__} and Ascend 910 AI software package version {v} does not "
                            "match, reference to the match info on: https://www.mindspore.cn/install")
 
+    def check_deps_version(self):
+        # in order to update the change of 'LD_LIBRARY_PATH' env, run a sub process
+        input_args = ["--mindspore_version=" + __version__]
+        for v in self.version:
+            input_args.append("--supported_version=" + v)
+        deps_version_checker = os.path.join(os.path.split(os.path.realpath(__file__))[0], "_check_deps_version.py")
+        call_cmd = [sys.executable, deps_version_checker] + input_args
+        process = subprocess.run(call_cmd, timeout=300, text=True, capture_output=True, check=False)
+        if process.stdout.strip() != "":
+            logger.warning(process.stdout.strip())
+
     def set_env(self):
         if not self.tbe_path:
             self._check_env()
@@ -198,12 +210,15 @@ class AscendEnvChecker(EnvChecker):
             import te
         except RuntimeError:
             if Path(self.tbe_path).is_dir():
-                os.environ['LD_LIBRARY_PATH'] = self.tbe_path
+                os.environ['LD_LIBRARY_PATH'] = self.tbe_path + ":" + os.environ['LD_LIBRARY_PATH']
             else:
                 raise EnvironmentError(
                     f"No such directory: {self.tbe_path}, Please check if Ascend 910 AI software package is "
                     "installed correctly.")
 
+        # check te version after set te env
+        self.check_deps_version()
+
         if Path(self.op_impl_path).is_dir():
             sys.path.append(self.op_impl_path)
         else: