update

minmingzhu · Nov 1, 2024 · 7ccf98d · 7ccf98d
1 parent 0971db6
commit 7ccf98d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/mllib-dal/src/main/native/DecisionForestRegressorImpl.cpp b/mllib-dal/src/main/native/DecisionForestRegressorImpl.cpp
@@ -391,7 +391,7 @@ Java_com_intel_oap_mllib_regression_RandomForestRegressorDALImpl_cRFRegressorTra
         t1 = std::chrono::high_resolution_clock::now();
         auto comm =
             preview::spmd::make_communicator<preview::spmd::backend::ccl>(
-                queue, executorNum, rank, kvs);
+                queue, executorNum, rank, kvs);:queue
         t2 = std::chrono::high_resolution_clock::now();
         duration =
             (float)std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1)

diff --git a/...dal/src/main/scala/com/intel/oap/mllib/classification/RandomForestClassifierDALImpl.scala b/...dal/src/main/scala/com/intel/oap/mllib/classification/RandomForestClassifierDALImpl.scala
@@ -76,6 +76,21 @@ class RandomForestClassifierDALImpl(val uid: String,
         "Please run on GPU device.")
     }
     rfcTimer.record("Data Convertion")
+
+    labeledPointsTables.mapPartitionsWithIndex { (rank, iter) =>
+      logInfo(s"set ZE_AFFINITY_MASK")
+      val gpuIndices = if (useDevice == "GPU") {
+        val resources = TaskContext.get().resources()
+        resources("gpu").addresses.map(_.toInt)
+      } else {
+        null
+      }
+      logInfo(s"set ZE_AFFINITY_MASK rank is $rank.")
+      logInfo(s"gpuIndices is ${gpuIndices.mkString(", ")}.")
+      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      Iterator.empty
+    }.count()
+
     val kvsIPPort = getOneCCLIPPort(labeledPointsTables)
     val training_breakdown_name = "RFClassifier_training_breakdown_" + executorNum;
 

diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/RandomForestRegressorDALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/RandomForestRegressorDALImpl.scala
@@ -70,6 +70,20 @@ class RandomForestRegressorDALImpl(val uid: String,
     }
     rfrTimer.record("Data Convertion")
 
+    labeledPointsTables.mapPartitionsWithIndex { (rank, iter) =>
+      logInfo(s"set ZE_AFFINITY_MASK")
+      val gpuIndices = if (useDevice == "GPU") {
+        val resources = TaskContext.get().resources()
+        resources("gpu").addresses.map(_.toInt)
+      } else {
+        null
+      }
+      logInfo(s"set ZE_AFFINITY_MASK rank is $rank.")
+      logInfo(s"gpuIndices is ${gpuIndices.mkString(", ")}.")
+      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      Iterator.empty
+    }.count()
+
     val kvsIPPort = getOneCCLIPPort(labeledPointsTables)
     val training_breakdown_name = "RFRegressor_training_breakdown_" + executorNum;