support throttle to fix ninja c1 oom (#8953)

* throttle * use throttle and fix ninja c1 * add deps * add server and with-cuda flag * enhance * portalocker * timeout 10 * update deps * rm * fix * fix * fix * auto format by CI * Update python/oneflow/test_utils/throttle.py Co-authored-by: Shenghang Tsai <jackalcooper@gmail.com> * auto format by CI * fix * fix * fix * rename and stop using shell * fix * fix * minor refactor * minor refactor * fix * fix Co-authored-by: jackalcooper <jackalcooper@gmail.com> Co-authored-by: oneflow-ci-bot <ci-bot@oneflow.org>
Oneflow-Inc · Aug 20, 2022 · 8f5745b · 8f5745b
1 parent cbea1ad
commit 8f5745b
Show file tree

Hide file tree

Showing 17 changed files with 76 additions and 14 deletions.
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -14,3 +14,4 @@ dataclasses; python_version<"3.7"
 cmakelang==0.6.13
 pytest-xdist
 rich
+portalocker
diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_batchnorm_relu.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_batchnorm_relu.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_bias_add.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_bias_add.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_conv.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_conv.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_conv2d_maxpool2d.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_conv2d_maxpool2d.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_conv_relu_add.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_conv_relu_add.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_lenet.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_lenet.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_maxpool_2d.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_maxpool_2d.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_resnet.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_resnet.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 import unittest
 import numpy as np

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_transpose_eliminate.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_nhwc_transpose_eliminate.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK-NOT: oneflow.transpose
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/auto_nhwc/test_resnet101_benchmark.py b/oneflow/ir/test/OneFlow/auto_nhwc/test_resnet101_benchmark.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 import unittest
 import numpy as np

diff --git a/oneflow/ir/test/OneFlow/cuda_code_gen/test_fuser_cast_scale.py b/oneflow/ir/test/OneFlow/cuda_code_gen/test_fuser_cast_scale.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: jit
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/folding/test_simple_multiply.py b/oneflow/ir/test/OneFlow/folding/test_simple_multiply.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK-NOT: oneflow.broadcast_mul
 
 import os

diff --git a/oneflow/ir/test/OneFlow/test_fuse_pad_conv.py b/oneflow/ir/test/OneFlow/test_fuse_pad_conv.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK-NOT: oneflow.pad
 
 import unittest

diff --git a/oneflow/ir/test/OneFlow/with_cuda/test_conv_bn_auto_nhwc.py b/oneflow/ir/test/OneFlow/with_cuda/test_conv_bn_auto_nhwc.py
@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# RUN: python3 %s | FileCheck %s
+# RUN: python3 -m oneflow.test_utils.throttle --with-cuda=%with_cuda python3 %s | FileCheck %s
 # CHECK: oneflow.transpose
 
 import os

diff --git a/oneflow/ir/test/lit.cfg.py b/oneflow/ir/test/lit.cfg.py
@@ -100,6 +100,7 @@
 tools = ["oneflow-opt", "oneflow-translate", "oneflow-runner"]
 tools.extend(
     [
+        ToolSubst("%with_cuda", config.BUILD_CUDA, unresolved="ignore"),
         ToolSubst("%linalg_test_lib_dir", config.llvm_lib_dir, unresolved="ignore"),
         ToolSubst("%test_exec_root", config.test_exec_root, unresolved="ignore"),
     ]

diff --git a/python/oneflow/test_utils/throttle.py b/python/oneflow/test_utils/throttle.py
@@ -0,0 +1,60 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+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.
+"""
+import argparse
+import hashlib
+import subprocess
+import portalocker
+import os
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Control when the script runs through special variables."
+    )
+    parser.add_argument(
+        "--with-cuda", type=int, default=1, help="whether has cuda device."
+    )
+    parser.add_argument("cmd", type=str, nargs="...", help="command to run")
+    return parser.parse_args()
+
+
+def hash_cli2gpu(cmd: list):
+    import pynvml
+
+    pynvml.nvmlInit()
+    slot = pynvml.nvmlDeviceGetCount()
+    hash = hashlib.sha1(" ".join(cmd).encode("utf-8")).hexdigest()
+    gpu_id = int(hash, 16) % slot
+    return [gpu_id]
+
+
+def main():
+    args = parse_args()
+    if args.with_cuda:
+        cuda_visible_devices = [str(i) for i in hash_cli2gpu(args.cmd)]
+        with portalocker.Lock(
+            ".oneflow-throttle-gpu-" + "-".join(cuda_visible_devices) + ".lock",
+            timeout=400,
+        ):
+            env = dict(os.environ, CUDA_VISIBLE_DEVICES=",".join(cuda_visible_devices))
+            return subprocess.call(args.cmd, env=env)
+    else:
+        return subprocess.call(args.cmd)
+
+
+if __name__ == "__main__":
+    returncode = main()
+    exit(returncode)