ditorch runs pytorch test script

ditorch/common_adapter.py

@@ -1,41 +0,0 @@
-import torch
-import torch.testing._internal.common_utils as common_utils
-
-
-class CudaNonDefaultStream:
-    def __enter__(self):
-        # Before starting CUDA test save currently active streams on all
-        # CUDA devices and set new non default streams to all CUDA devices
-        # to ensure CUDA tests do not use default stream by mistake.
-        beforeDevice = torch.cuda.current_device()
-        self.beforeStreams = []
-        for d in range(torch.cuda.device_count()):
-            self.beforeStreams.append(torch.cuda.current_stream(d))
-            deviceStream = torch.cuda.Stream(device=d)
-            self.beforeStreams[-1].synchronize()
-            """
-            torch._C._cuda_setStream(stream_id=deviceStream.stream_id,
-                                     device_index=deviceStream.device_index,
-                                     device_type=deviceStream.device_type)
-            """
-            torch.cuda.set_stream(deviceStream)
-
-        # torch._C._cuda_setDevice(beforeDevice)
-        torch.cuda.set_device(beforeDevice)
-
-    def __exit__(self, exec_type, exec_value, traceback):
-        # After completing CUDA test load previously active streams on all
-        # CUDA devices.
-        beforeDevice = torch.cuda.current_device()
-        for d in range(torch.cuda.device_count()):
-            """
-            torch._C._cuda_setStream(stream_id=self.beforeStreams[d].stream_id,
-                                     device_index=self.beforeStreams[d].device_index,
-                                     device_type=self.beforeStreams[d].device_type)
-            """
-            torch.cuda.set_stream(self.beforeStreams[d])
-        # torch._C._cuda_setDevice(beforeDevice)
-        torch.cuda.set_device(beforeDevice)
-
-
-common_utils.CudaNonDefaultStream = CudaNonDefaultStream

ditorch/test/discover_pytorch_test_case.py

@@ -0,0 +1,62 @@
+import unittest
+import ditorch  # noqa: F401
+import sys
+import json
+
+
+def discover_test_cases_recursively(suite_or_case):
+    if isinstance(suite_or_case, unittest.TestCase):
+        return [suite_or_case]
+    rc = []
+    for element in suite_or_case:
+        rc.extend(discover_test_cases_recursively(element))
+    return rc
+
+
+def get_test_full_names(test_cases):
+    full_name_list = []
+    for case in test_cases:
+        id = case.id()
+        module_name = id.split(".")[0] + ".py"
+        test_name = id[id.find(".") + 1 :]
+        full_name = module_name + " " + test_name
+        full_name_list.append(full_name)
+    return full_name_list
+
+
+def discover_all_test_case(path="."):
+    loader = unittest.TestLoader()
+    test_cases = loader.discover(path)
+    test_cases = discover_test_cases_recursively(test_cases)
+    return test_cases
+
+
+def dump_all_test_case_id_to_file(test_cases, path="all_test_cases.json"):
+    test_case_ids = {}
+    case_num = 0
+    for case in test_cases:
+        case_id = case.id()
+        print(case_id)
+        case_num += 1
+        module_name = case_id.split(".")[0] + ".py"
+        test_name = case_id[case_id.find(".") + 1 :]
+        if not module_name.startswith("test_"):
+            continue
+
+        if module_name not in test_case_ids:
+            test_case_ids[module_name] = [test_name]
+        else:
+            test_case_ids[module_name].append(test_name)
+
+    with open(path, "wt") as f:
+        json.dump(test_case_ids, f)
+
+    print(f"dumped {case_num} test cases from {len(test_case_ids)} files to {path}")
+    return test_case_ids
+
+
+if __name__ == "__main__":
+    print(f"discover:{sys.argv}")
+    path = sys.argv[1] if len(sys.argv) > 1 else "."
+    all_tests_case = discover_all_test_case(path)
+    dump_all_test_case_id_to_file(all_tests_case)

ditorch/test/test_torch_testcase.py

@@ -1,21 +1,30 @@
 import os
+import subprocess
+import json
 
 
-def run_test(file_path):
-    os.makedirs("temp", exist_ok=True)
+def copy_add_ditorch_import_to_pytorch_test_file(file_name, pytorch_dir, dest_dir="temp"):
+    file_path = f"{pytorch_dir}/test/" + file_name
+    os.makedirs(dest_dir, exist_ok=True)
+    # We only add a line "import ditorch" to the original pytorch test script, and make no changes other than that.
     with open(file_path, "rt") as torch_test_source_script_file:
         content = torch_test_source_script_file.read()
         content = "import ditorch\n" + content
-    new_file_name = "temp/" + file_path.split("/")[-1]
-    pass
+    new_file_name = dest_dir + "/" + file_name
     with open(new_file_name, "w") as new_file:
         new_file.write(content)
-    print(f"test {file_path} over")
 
 
-TORCH_TEST_SCRIPT_FILE = [
-    "test_ops.py",
-]
+def run_command_in_sub_process(commands):
+    print(commands)
+    result = subprocess.run(commands, shell=True, text=True, capture_output=True)
+    print(result.stdout)
+    print(result.stderr)
+    if result.returncode != 0:
+        print(f"Run {commands} FAILED")
+    else:
+        print(f"Run {commands} PASSED")
+    print("\n\n\n")
 
 
 def main():
@@ -26,9 +35,24 @@ def main():
     if not os.path.isdir(pytorch_dir):
         print(f"{pytorch_dir} is not exist")
         return -1
-    for file_path in TORCH_TEST_SCRIPT_FILE:
-        full_path = pytorch_dir + "/test/" + file_path
-        run_test(full_path)
+
+    # all_test_case = discover_all_test_case(pytorch_dir + "/test/")
+
+    # Finding runnable test cases needs to be done in a process where the device is available.
+    # The current process is a clean process that does not use device(import ditorch).
+    run_command_in_sub_process(f"python ditorch/test/discover_pytorch_test_case.py {pytorch_dir}/test")
+    with open("all_test_cases.json", "r") as f:
+        test_case_ids = json.load(f)
+
+    source_files = test_case_ids.keys()
+    for file_name in source_files:
+        copy_add_ditorch_import_to_pytorch_test_file(file_name, pytorch_dir, dest_dir="temp")
+
+    for test_script_file, test_cases in test_case_ids.items():
+        for case in test_cases:
+            commands = f"python temp/{test_script_file} {case} -v --subprocess"
+            print(f"Running {commands}")
+            run_command_in_sub_process(commands)
 
 
 if __name__ == "__main__":

ditorch/torch_npu_adapter.py

@@ -18,7 +18,26 @@ def current_stream(device=None):
 
 
 def get_device_capability():
-    return (7, 5)
+    return (8, 0)
 
 
 torch.cuda.get_device_capability = get_device_capability
+
+torch._C._cuda_setStream = torch_npu._C._npu_setStream
+
+torch._C._cuda_setDevice = torch_npu._C._npu_setDevice
+
+
+def _set_cached_tensors_enabled(enable):
+    print("_set_cached_tensors_enabled not supported in torch_npu")
+
+
+torch._C._set_cached_tensors_enabled = _set_cached_tensors_enabled
+
+
+def _add_cached_tensor(tensor):
+    print("_add_cached_tensor not supported in torch_npu")
+    raise NotImplementedError
+
+
+torch._C._add_cached_tensor = _add_cached_tensor