What is the earliest way to call a Java/Scala method?

Viewed 231

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.

2 Answers

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?

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
  }
}
Related