5

I use some dirty code to set environment variable in Scala from second answer of this question

I test this in IDE(IDEA Intellij) and set OMP_NUM_THREADS at the beginning of my class.

import org.scalatest.{FlatSpec, Matchers}
class MyTest extends FlatSpec with Matchers {
  val m = Map("OMP_NUM_THREADS" -> "1")
  EnvHacker.setEnv(m)

After set, I could read from System.env, it works. But when my program runs, it does not use this. I tried set it in static block, but still not work.

But if I set it in IDE run configuration(before JVM run), it works and runs as I expect. So seems it is read before I modify the variable.

Or in other word, I have a piece of code, what is the earliest way to call it in Java/Scala. e.g. static block is called before the first line of main method.

Some details updated:

I am using tensorflow-mkl Java API, it would read System environment variable OMP_NUM_THREADS at some time, according to my test result, this operation is before the system static block. However, I want to control in code because I do not know the configuration expected without code logic.

Litchy
  • 623
  • 7
  • 23
  • 1
    How does your program retrieves this environment variable? Could it be that it retrieves it before you run your `EnvHacker.setEnv`? – Gaël J May 17 '21 at 08:10
  • @GaëlJ I think so, so is there any condition that it retrieves before the first line of running code? – Litchy May 17 '21 at 08:21
  • 1
    It all depends on where/when your program retrieves this environment variable. If this is part of a framework/library you use, either they provided some kind of "hooks" which you can call before their code run, or it's probably a dead end. – Gaël J May 17 '21 at 08:27
  • To clarify, do you need this environment variable only in some tests or at runtime of the app? – Gaël J May 17 '21 at 08:28
  • at run time, I want to do some more control of process, but do not want to create other process at this time. If no way, maybe I would create another using ProcessBuilder – Litchy May 17 '21 at 08:31
  • @GaëlJG I tested to set the environment variable at static block(before code execution), but still does not work. – Litchy May 18 '21 at 07:10
  • @GaëlJ Thanks for your answer, I have simplified the question according to your help. – Litchy May 18 '21 at 07:27
  • You sample code is specific to test case with scalatest. Do you want a solution which works with scalatest ? Solution to this problem is highly dependent on how you are "running" your code. Scalatest provides with [beforeAndAfter](https://www.scalatest.org/user_guide/sharing_fixtures#beforeAndAfter) hook for such purposes. – sarveshseri May 25 '21 at 14:57

2 Answers2

1

I think this is an XY problem: you want to control OpenMP via "OMP_NUM_THREADS" env but you can't or don't want to set the actual environment variables on your process for some reason? Is that right?

How are you invoking OpenMP? OpenMP is a C lib https://www.openmp.org/ , so it won't notice any changes you make to the java.lang.ProcessEnvironment#theEnvironment field, no matter how early you set it.

If you are invoking OpenMP via an exec call, then you should be able to pass a new environment to it.

If you are invoking OpenMP via JNI, then you won't be able to change your environment variables from Java, you will need to actually set the env when you start the process. See Can I set an Environment Variable for Java Native Interface (JNI) libraries?

Could you instead use omp_set_num_threads() rather than OMP_NUM_THREADS env?

Rich
  • 15,048
  • 2
  • 66
  • 119
  • Exactly! This is an XY problem. I add some details of the question. You said "you won't be able to change your environment variables from Java", but the hacking code seems to work. So that is turns out tensorflow-mkl java api does not invoke OpenMP from JNI? Am I getting it right? – Litchy May 31 '21 at 06:43
  • I don't know how tensorflow-mkl java invokes OpenMP and I can't figure it out from 5 mins of googling. I think your real question is more like "how do I control the number of threads used by OpenMP inside tensorflow-mkl java", but I'm not very familiar with that toolchain. I think in an ideal world you would close this question and open a new question with that title and with some details of your setup, but if you have a solution that works for you then maybe that is not needed. This question won't be useful to any other Java devs as it has a misleading title IMO. GL! – Rich Jun 07 '21 at 10:09
-1

I think this might do what you want, I borrowed the env hacker (very nasty hack!) from the other question, my understanding is you want to set some environment variables and have them available in the main method (essentially as early as possible).

import java.util.{Collections, Map => JavaMap}
import scala.collection.JavaConverters._

object EnvHacker {
    /**
     * Portable method for setting env vars on both *nix and Windows.
     * @see http://stackoverflow.com/a/7201825/293064
     */
    def setEnv(newEnv: Map[String, String]): Unit = {
        try {
            val processEnvironmentClass = Class.forName("java.lang.ProcessEnvironment")
            val theEnvironmentField = processEnvironmentClass.getDeclaredField("theEnvironment")
            theEnvironmentField.setAccessible(true)
            val env = theEnvironmentField.get(null).asInstanceOf[JavaMap[String, String]]
            env.putAll(newEnv.asJava)
            val theCaseInsensitiveEnvironmentField = processEnvironmentClass.getDeclaredField("theCaseInsensitiveEnvironment")
            theCaseInsensitiveEnvironmentField.setAccessible(true)
            val cienv = theCaseInsensitiveEnvironmentField.get(null).asInstanceOf[JavaMap[String, String]]
            cienv.putAll(newEnv.asJava)
        } catch {
            case e: NoSuchFieldException =>
                try {
                    val classes = classOf[Collections].getDeclaredClasses
                    val env = System.getenv()
                    for {cl <- classes} {
                        if (cl.getName == "java.util.Collections$UnmodifiableMap") {
                            val field = cl.getDeclaredField("m")
                            field.setAccessible(true)
                            val obj = field.get(env)
                            val map = obj.asInstanceOf[JavaMap[String, String]]
                            map.clear()
                            map.putAll(newEnv.asJava)
                        }
                    }
                } catch {
                    case e2: Exception => e2.printStackTrace()
                }

            case e1: Exception => e1.printStackTrace()
        }
    }
}


class Main extends {
  val NumThreads = "OP_NUM_THREADS"
  EnvHacker.setEnv(Map(NumThreads -> "100"))
}

object Main extends Main {
  def main(args: Array[String]): Unit = {
    println(System.getenv(NumThreads)) // prints 100
  }
}
Ahmad Ragab
  • 1,067
  • 1
  • 12
  • 24