update

oap-project · Sep 10, 2024 · 51f03c8 · 51f03c8
1 parent 04731dd
commit 51f03c8
Show file tree

Hide file tree

Showing 8 changed files with 10 additions and 10 deletions.
diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/OneCCL.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/OneCCL.scala
@@ -39,9 +39,9 @@ object OneCCL extends Logging {
       s"commSize, ${cclParam.getCommSize}, rankId: ${cclParam.getRankId}")
   }
 
-  // Run on Executor
-  def setExecutorEnv(key: String, value: String): Unit = {
-    setEnv(key, value)
+  // Set specified by values for each Executor
+  def setAffinityMask(rankId: String): Unit = {
+      setEnv("ZE_AFFINITY_MASK", rankId)
   }
 
   // Run on Executor

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
@@ -82,7 +82,7 @@ class RandomForestClassifierDALImpl(val uid: String,
       } else {
         null
       }
-      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      OneCCL.setAffinityMask(gpuIndices(0).toString())
       Iterator.empty
     }.count()
 

diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/clustering/KMeansDALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/clustering/KMeansDALImpl.scala
@@ -58,7 +58,7 @@ class KMeansDALImpl(var nClusters: Int,
       } else {
         null
       }
-      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      OneCCL.setAffinityMask(gpuIndices(0).toString())
       Iterator.empty
     }.count()
 

diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/feature/PCADALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/feature/PCADALImpl.scala
@@ -66,7 +66,7 @@ class PCADALImpl(val k: Int,
       } else {
         null
       }
-      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      OneCCL.setAffinityMask(gpuIndices(0).toString())
       Iterator.empty
     }.count()
 

diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/LinearRegressionDALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/LinearRegressionDALImpl.scala
@@ -113,7 +113,7 @@ class LinearRegressionDALImpl( val fitIntercept: Boolean,
       } else {
         null
       }
-      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      OneCCL.setAffinityMask(gpuIndices(0).toString())
       Iterator.empty
     }.count()
 

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
@@ -76,7 +76,7 @@ class RandomForestRegressorDALImpl(val uid: String,
       } else {
         null
       }
-      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      OneCCL.setAffinityMask(gpuIndices(0).toString())
       Iterator.empty
     }.count()
 

diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/stat/CorrelationDALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/stat/CorrelationDALImpl.scala
@@ -59,7 +59,7 @@ class CorrelationDALImpl(
       } else {
         null
       }
-      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      OneCCL.setAffinityMask(gpuIndices(0).toString())
       Iterator.empty
     }.count()
 

diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/stat/SummarizerDALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/stat/SummarizerDALImpl.scala
@@ -60,7 +60,7 @@ class SummarizerDALImpl(val executorNum: Int,
       } else {
         null
       }
-      OneCCL.setExecutorEnv("ZE_AFFINITY_MASK", gpuIndices(0).toString())
+      OneCCL.setAffinityMask(gpuIndices(0).toString())
       Iterator.empty
     }.count()