Refactor versions wheels (#126)

* fix * publishing * add
Oneflow-Inc · Nov 15, 2024 · 29bb4b2 · 29bb4b2
1 parent f1db695
commit 29bb4b2
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 21 deletions.
diff --git a/dist/index.js b/dist/index.js
diff --git a/src/buildOneFlow.ts b/src/buildOneFlow.ts
@@ -91,13 +91,13 @@ async function buildWithConda(): Promise<void> {
 
 const ProductionCommit = '60785746d8a0c1c3b16e094ad653d1a18c519c7f'
 const CUDA12ProductionCommit = 'bba403da5da597c3f7905dfc52b24989ee21ffb1'
-const V2ProductionCommit = '6455f9b8154333333e6285fde3747aaac4a92929'
+const V3ProductionCommit = '7725348ba96e391b4660a4190e7d4e757e7cb83f'
 
 const CUDA_120_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda12.0:${CUDA12ProductionCommit}`
-const CUDA_121_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda12.1:${V2ProductionCommit}`
-const CUDA_122_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda12.2:${V2ProductionCommit}`
+const CUDA_121_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda12.1:${V3ProductionCommit}`
+const CUDA_122_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda12.2:${V3ProductionCommit}`
 
-const CUDA_118_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda11.8:${V2ProductionCommit}`
+const CUDA_118_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda11.8:${V3ProductionCommit}`
 const CUDA_116_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda11.6:${ProductionCommit}`
 const CUDA_117_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda11.7:${ProductionCommit}`
 const CUDA_115_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda11.5:${ProductionCommit}`
@@ -107,7 +107,7 @@ const CUDA_112_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014
 const CUDA_110_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda11.0:${ProductionCommit}`
 const CUDA_102_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cuda10.2:${ProductionCommit}`
 
-const CUDA_CPU_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cpu:${V2ProductionCommit}`
+const CUDA_CPU_IMG_TAG = `registry.cn-beijing.aliyuncs.com/oneflow/manylinux2014_x86_64_cpu:${V3ProductionCommit}`
 
 type CudaVersion =
   | '10.2'

diff --git a/src/utils/docker.ts b/src/utils/docker.ts
@@ -182,21 +182,21 @@ async function buildAndMakeWheel(
   let postProcessCmds = [runCPack(container, buildDir)]
   let nvLibs: string[] = []
   if (useNVWheels) {
-    nvLibs = [
-      'libcudnn_cnn_infer.so.8',
-      'libcudnn_cnn_train.so.8',
-      'libcudnn_ops_infer.so.8',
-      'libcudnn_ops_train.so.8',
-      'libcublas.so.11',
-      'libcublasLt.so.11',
-      'libcublas.so.12',
-      'libcublasLt.so.12',
-      'libnccl.so.2',
-      'libcufft.so.11',
-      'libcufft.so.12',
-      'libcusparse.so.11',
-      'libcusparse.so.12'
-    ]
+    for (let suffix = 0; suffix <= 20; suffix++) {
+      nvLibs = nvLibs.concat([
+        `libcudnn_cnn_infer.so.${suffix}`,
+        `libcudnn_ops_infer.so.${suffix}`,
+        `libcublas.so.${suffix}`,
+        `libcublasLt.so.${suffix}`,
+        `libnccl.so.${suffix}`,
+        `libcufft.so.${suffix}`,
+        `libcusparse.so.${suffix}`,
+        `libcupti.so.${suffix}`,
+        `libcurand.so.${suffix}`,
+        `libcusolver.so.${suffix}`,
+        `libnvJitLink.so.${suffix}`
+      ])
+    }
   }
   const nvLibsExcludes = Array.prototype.concat.apply(
     [],