Implement [-t|--threads] command line argument for threads

NEWS.md

@@ -72,6 +72,15 @@ Language changes
 
 * Color now defaults to on when stdout and stderr are TTYs ([#34347])
 
+Command-line option changes
+---------------------------
+
+  * `-t N`, `--threads N` starts Julia with `N` threads. This option takes precedence over
+    `JULIA_NUM_THREADS`. The specified number of threads also propagates to worker
+    processes spawned using the `-p`/`--procs` or `--machine-file` command line arguments.
+    In order to set number of threads for worker processes spawned with `addprocs` use the
+    `exeflags` keyword argument, e.g. `` addprocs(...; exeflags=`--threads 4`) `` ([#35108]).
+
 Multi-threading changes
 -----------------------
 

base/options.jl

@@ -9,6 +9,7 @@ struct JLOptions
     commands::Ptr{Ptr{UInt8}} # (e)eval, (E)print, (L)load
     image_file::Ptr{UInt8}
     cpu_target::Ptr{UInt8}
+    nthreads::Int32
     nprocs::Int32
     machine_file::Ptr{UInt8}
     project::Ptr{UInt8}

doc/man/julia.1

@@ -101,6 +101,10 @@ Evaluate <expr> and display the result
 -L, --load <file>
 Load <file> immediately on all processors
 
+.TP
+-t, --threads <n>
+Enable n threads
+
 .TP
 -p, --procs <n>
 Run n local processes

doc/src/manual/environment-variables.md

@@ -182,7 +182,7 @@ A [`Float64`](@ref) that sets the value of `Distributed.worker_timeout()` (defau
 This function gives the number of seconds a worker process will wait for
 a master process to establish a connection before dying.
 
-### `JULIA_NUM_THREADS`
+### [`JULIA_NUM_THREADS`](@id JULIA_NUM_THREADS)
 
 An unsigned 64-bit integer (`uint64_t`) that sets the maximum number of threads
 available to Julia. If `$JULIA_NUM_THREADS` exceeds the number of available
@@ -195,6 +195,10 @@ set to `1`.
 
     `JULIA_NUM_THREADS` must be defined before starting julia; defining it in `startup.jl` is too late in the startup process.
 
+!!! compat "Julia 1.5"
+    In Julia 1.5 and above the number of threads can also be specified on startup
+    using the `-t`/`--threads` command line argument.
+
 ### `JULIA_THREAD_SLEEP_THRESHOLD`
 
 If set to a string that starts with the case-insensitive substring `"infinite"`,

doc/src/manual/getting-started.md

@@ -109,6 +109,7 @@ julia [switches] -- [programfile] [args...]
 |`-e`, `--eval <expr>`                  |Evaluate `<expr>`|
 |`-E`, `--print <expr>`                 |Evaluate `<expr>` and display the result|
 |`-L`, `--load <file>`                  |Load `<file>` immediately on all processors|
+|`-t`, `--threads {N\|auto`}            |Enable N threads; `auto` currently sets N to the number of local CPU threads but this might change in the future|
 |`-p`, `--procs {N\|auto`}              |Integer value N launches N additional local worker processes; `auto` launches as many workers as the number of local CPU threads (logical cores)|
 |`--machine-file <file>`                |Run processes on hosts listed in `<file>`|
 |`-i`                                   |Interactive mode; REPL runs and `isinteractive()` is true|

doc/src/manual/parallel-computing.md

@@ -231,28 +231,21 @@ julia> Threads.nthreads()
 1
 ```
 
-The number of threads Julia starts up with is controlled by an environment variable called `JULIA_NUM_THREADS`.
-Now, let's start up Julia with 4 threads:
+The number of threads Julia starts up with is controlled either by using the
+`-t`/`--threads` command line argument or by using the
+[`JULIA_NUM_THREADS`](@ref JULIA_NUM_THREADS) environment variable. When both are
+specified, then `-t`/`--threads` takes precedence.
 
-Bash on Linux/OSX:
+!!! compat "Julia 1.5"
+    The `-t`/`--threads` command line argument requires at least Julia 1.5.
+    In older versions you must use the environment variable instead.
 
-```bash
-export JULIA_NUM_THREADS=4
-```
-
-C shell on Linux/OSX, CMD on Windows:
+Lets start Julia with 4 threads:
 
 ```bash
-set JULIA_NUM_THREADS=4
+$ julia --threads 4
 ```
 
-Powershell on Windows:
-
-```powershell
-$env:JULIA_NUM_THREADS=4
-```
-
-
 Let's verify there are 4 threads at our disposal.
 
 ```julia-repl
@@ -267,6 +260,29 @@ julia> Threads.threadid()
 1
 ```
 
+!!! note
+    If you prefer to use the environment variable you can set it as follows in
+    Bash (Linux/macOS):
+    ```bash
+    export JULIA_NUM_THREADS=4
+    ```
+    C shell on Linux/macOS, CMD on Windows:
+    ```bash
+    set JULIA_NUM_THREADS=4
+    ```
+    Powershell on Windows:
+    ```powershell
+    $env:JULIA_NUM_THREADS=4
+    ```
+    Note that this must be done *before* starting Julia.
+
+!!! note
+    The number of threads specified with `-t`/`--threads` is propagated to worker processes
+    that are spawned using the `-p`/`--procs` or `--machine-file` command line options.
+    For example, `julia -p2 -t2` spawns 1 main process with 2 worker processes, and all
+    three processes have 2 threads enabled. For more fine grained control over worker
+    threads use [`addprocs`](@ref) and pass `-t`/`--threads` as `exeflags`.
+
 ## The `@threads` Macro
 
 Let's work a simple example using our native threads. Let us create an array of zeros:

src/jloptions.c

@@ -35,6 +35,7 @@ jl_options_t jl_options = { 0,    // quiet
                             NULL, // cmds
                             NULL, // image_file (will be filled in below)
                             NULL, // cpu_target ("native", "core2", etc...)
+                            0,    // nthreads
                             0,    // nprocs
                             NULL, // machine_file
                             NULL, // project
@@ -97,6 +98,9 @@ static const char opts[]  =
     " -L, --load <file>         Load <file> immediately on all processors\n\n"
 
     // parallel options
+    " -t, --threads {N|auto}    Enable N threads; \"auto\" currently sets N to the number of local\n"
+    "                           CPU threads but this might change in the future\n"
+    " -t, --threads {N|auto}    Enable N threads. \"auto\" sets N to the number of local CPU threads.\n"
     " -p, --procs {N|auto}      Integer value N launches N additional local worker processes\n"
     "                           \"auto\" launches as many workers as the number of local CPU threads (logical cores)\n"
     " --machine-file <file>     Run processes on hosts listed in <file>\n\n"
@@ -190,7 +194,7 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
            opt_machine_file,
            opt_project,
     };
-    static const char* const shortopts = "+vhqH:e:E:L:J:C:ip:O:g:";
+    static const char* const shortopts = "+vhqH:e:E:L:J:C:it:p:O:g:";
     static const struct option longopts[] = {
         // exposed command line options
         // NOTE: This set of required arguments need to be kept in sync
@@ -209,6 +213,7 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
         { "compiled-modules",    required_argument, 0, opt_compiled_modules },
         { "cpu-target",      required_argument, 0, 'C' },
         { "procs",           required_argument, 0, 'p' },
+        { "threads",         required_argument, 0, 't' },
         { "machine-file",    required_argument, 0, opt_machine_file },
         { "project",         optional_argument, 0, opt_project },
         { "color",           required_argument, 0, opt_color },
@@ -388,6 +393,18 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
             if (!jl_options.cpu_target)
                 jl_error("julia: failed to allocate memory");
             break;
+        case 't': // threads
+            errno = 0;
+            if (!strcmp(optarg,"auto")) {
+                jl_options.nthreads = -1;
+            }
+            else {
+                long nthreads = strtol(optarg, &endptr, 10);
+                if (errno != 0 || optarg == endptr || *endptr != 0 || nthreads < 1 || nthreads >= INT_MAX)
+                    jl_errorf("julia: -t,--threads=<n> must be an integer >= 1");
+                jl_options.nthreads = (int)nthreads;
+            }
+            break;
         case 'p': // procs
             errno = 0;
             if (!strcmp(optarg,"auto")) {

src/julia.h

@@ -1900,6 +1900,7 @@ typedef struct {
     const char **cmds;
     const char *image_file;
     const char *cpu_target;
+    int32_t nthreads;
     int32_t nprocs;
     const char *machine_file;
     const char *project;

src/threading.c

@@ -407,8 +407,11 @@ void jl_init_threading(void)
     // how many threads available, usable
     int max_threads = jl_cpu_threads();
     jl_n_threads = JULIA_NUM_THREADS;
-    cp = getenv(NUM_THREADS_NAME);
-    if (cp)
+    if (jl_options.nthreads < 0) // --threads=auto
+        jl_n_threads = max_threads;
+    else if (jl_options.nthreads > 0) // --threads=N
+        jl_n_threads = jl_options.nthreads;
+    else if ((cp = getenv(NUM_THREADS_NAME)))
         jl_n_threads = (uint64_t)strtol(cp, NULL, 10);
     if (jl_n_threads > max_threads)
         jl_n_threads = max_threads;

stdlib/Distributed/src/cluster.jl

@@ -1298,6 +1298,7 @@ end
 
 write_cookie(io::IO) = print(io.in, string(cluster_cookie(), "\n"))
 
+# Starts workers specified by (-n|--procs) and --machine-file command line options
 function process_opts(opts)
     # startup worker.
     # opts.startupfile, opts.load, etc should should not be processed for workers.
@@ -1310,14 +1311,17 @@ function process_opts(opts)
         end
     end
 
+    # Propagate --threads to workers
+    exeflags = opts.nthreads > 0 ? `--threads=$(opts.nthreads)` : ``
+
     # add processors
     if opts.nprocs > 0
-        addprocs(opts.nprocs)
+        addprocs(opts.nprocs; exeflags=exeflags)
     end
 
     # load processes from machine file
     if opts.machine_file != C_NULL
-        addprocs(load_machine_file(unsafe_string(opts.machine_file)))
+        addprocs(load_machine_file(unsafe_string(opts.machine_file)); exeflags=exeflags)
     end
     return nothing
 end

test/cmdlineargs.jl

@@ -186,6 +186,32 @@ let exename = `$(Base.julia_cmd()) --startup-file=no`
     @test !success(`$exename -C invalidtarget`)
     @test !success(`$exename --cpu-target=invalidtarget`)
 
+    # -t, --threads
+    code = "print(Threads.nthreads())"
+    cpu_threads = ccall(:jl_cpu_threads, Int32, ())
+    @test string(cpu_threads) ==
+          read(`$exename --threads auto -e $code`, String) ==
+          read(`$exename --threads=auto -e $code`, String) ==
+          read(`$exename -tauto -e $code`, String) ==
+          read(`$exename -t auto -e $code`, String) ==
+          read(`$exename -t $(cpu_threads+1) -e $code`, String)
+    if cpu_threads > 1
+        for nt in (nothing, "1"); withenv("JULIA_NUM_THREADS"=>nt) do
+            @test read(`$exename --threads 2 -e $code`, String) ==
+                  read(`$exename --threads=2 -e $code`, String) ==
+                  read(`$exename -t2 -e $code`, String) ==
+                  read(`$exename -t 2 -e $code`, String) == "2"
+        end end
+    end
+    @test !success(`$exename -t 0`)
+    @test !success(`$exename -t -1`)
+
+    # Combining --threads and --procs: --threads does propagate
+    if cpu_threads > 1; withenv("JULIA_NUM_THREADS"=>nothing) do
+        code = "print(sum(remotecall_fetch(Threads.nthreads, x) for x in procs()))"
+        @test read(`$exename -p2 -t2 -e $code`, String) == "6"
+    end end
+
     # --procs
     @test readchomp(`$exename -q -p 2 -e "println(nworkers())"`) == "2"
     @test !success(`$exename -p 0`)