diff --git a/NNContext.scala b/NNContext.scala index 9d7f8d9c49e..7c0f41bef55 100644 --- a/NNContext.scala +++ b/NNContext.scala @@ -209,35 +209,39 @@ object NNContext { private[zoo] def initConf(zooConf: SparkConf) : Unit = { // check env and set spark conf // Set default value - var kmpAffinity = "granularity=fine,compact,1,0" - var kmpBlockTime = "0" - var kmpSettings = "1" - var ompNumThreads = "1" - // Set value with env - if (env.contains("KMP_AFFINITY")) { - kmpAffinity = env("KMP_AFFINITY") - } - if (env.contains("KMP_BLOCKTIME")) { - kmpBlockTime = env("KMP_BLOCKTIME") - } - if (env.contains("KMP_SETTINGS")) { - kmpSettings = env("KMP_SETTINGS") - } - if (env.contains("OMP_NUM_THREADS")) { - ompNumThreads = env("OMP_NUM_THREADS") - } else if (env.contains("ZOO_NUM_MKLTHREADS")) { - if (env("ZOO_NUM_MKLTHREADS").equalsIgnoreCase("all")) { - val cores = Runtime.getRuntime.availableProcessors() - ompNumThreads = cores.toString - } else { - ompNumThreads = env("ZOO_NUM_MKLTHREADS") + // We should skip this env, when engineType is mkldnn. + if (System.getProperty("bigdl.engineType", "mklblas") + .toLowerCase() == "mklblas") { + var kmpAffinity = "granularity=fine,compact,1,0" + var kmpBlockTime = "0" + var kmpSettings = "1" + var ompNumThreads = "1" + // Set value with env + if (env.contains("KMP_AFFINITY")) { + kmpAffinity = env("KMP_AFFINITY") + } + if (env.contains("KMP_BLOCKTIME")) { + kmpBlockTime = env("KMP_BLOCKTIME") + } + if (env.contains("KMP_SETTINGS")) { + kmpSettings = env("KMP_SETTINGS") + } + if (env.contains("OMP_NUM_THREADS")) { + ompNumThreads = env("OMP_NUM_THREADS") + } else if (env.contains("ZOO_NUM_MKLTHREADS")) { + if (env("ZOO_NUM_MKLTHREADS").equalsIgnoreCase("all")) { + val cores = Runtime.getRuntime.availableProcessors() + ompNumThreads = cores.toString + } else { + ompNumThreads = env("ZOO_NUM_MKLTHREADS") + } } + // Set Spark Conf + zooConf.setExecutorEnv("KMP_AFFINITY", kmpAffinity) + zooConf.setExecutorEnv("KMP_BLOCKTIME", kmpBlockTime) + zooConf.setExecutorEnv("KMP_SETTINGS", kmpSettings) + zooConf.setExecutorEnv("OMP_NUM_THREADS", ompNumThreads) } - // Set Spark Conf - zooConf.setExecutorEnv("KMP_AFFINITY", kmpAffinity) - zooConf.setExecutorEnv("KMP_BLOCKTIME", kmpBlockTime) - zooConf.setExecutorEnv("KMP_SETTINGS", kmpSettings) - zooConf.setExecutorEnv("OMP_NUM_THREADS", ompNumThreads) }