diff --git a/src/bindings/lua/wreck.lua b/src/bindings/lua/wreck.lua index 819342241aa5..944f3648498e 100644 --- a/src/bindings/lua/wreck.lua +++ b/src/bindings/lua/wreck.lua @@ -52,6 +52,7 @@ local default_opts = { ['error'] = { char = "E", arg = "FILENAME" }, ['input'] = { char = "i", arg = "HOW" }, ['label-io'] = { char = "l", }, + ['skip-env'] = { char = "S", }, ['options'] = { char = 'o', arg = "OPTIONS.." }, } @@ -119,6 +120,7 @@ function wreck:usage() -i /dev/null:* closes all stdin, etc.) -E, --error=FILENAME Send stderr to a different location than stdout. -l, --labelio Prefix lines of output with task id + -S, --skip-env Skip export of environment to job ]]) for _,v in pairs (self.extra_options) do local optstr = v.name .. (v.arg and "="..v.arg or "") @@ -131,7 +133,8 @@ function wreck:usage() end end -local function get_filtered_env () + +function wreck.get_filtered_env () local env = posix.getenv() env.HOSTNAME = nil env.ENVIRONMENT = nil @@ -142,6 +145,19 @@ local function get_filtered_env () return (env) end +local function get_job_env (arg) + local f = arg.flux + local env = wreck.get_filtered_env () + local default_env = {} + if f then default_env = f:kvs_get ("lwj.environ") or {} end + for k,v in pairs (env) do + -- If same key=value is already in default env no need to + -- export it again, remove: + if default_env[k] == env[k] then env[k] = nil end + end + return (env) +end + local function job_kvspath (f, id) assert (id, "Required argument id missing!") @@ -332,7 +348,7 @@ function wreck:jobreq () ntasks = self.ntasks, ncores = self.ncores, cmdline = self.cmdline, - environ = get_filtered_env (), + environ = self.opts.S and {} or get_job_env { flux = self.flux }, cwd = posix.getcwd (), walltime =self.walltime or 0, diff --git a/src/cmd/flux-wreck b/src/cmd/flux-wreck index 55a2f6f851a8..25fc790ac885 100755 --- a/src/cmd/flux-wreck +++ b/src/cmd/flux-wreck @@ -416,6 +416,66 @@ prog:SubCommand { end } +prog:SubCommand { + name = "setenv", + usage = "[VAR=VALUE|all]", + description = "Export environment to all jobs", + handler = function (self, arg) + local f = assert (require 'flux'.new ()) + local getenv = wreck.get_filtered_env + local env = f:kvs_get ("lwj.environ") or {} + for _,expr in ipairs (arg) do + if expr == "all" then + for k,v in pairs(getenv()) do + env[k] = v + end + else + local var,val = expr:match("(.*)=(.*)") + env[var] = val + end + end + f:kvs_put ("lwj.environ", env) + f:kvs_commit () + end +} + +prog:SubCommand { + name = "unsetenv", + usage = "VAR [VAR...]", + description = "Export environment to all jobs", + handler = function (self, arg) + local f = assert (require 'flux'.new ()) + local env = f:kvs_get ("lwj.environ") or {} + for _,var in ipairs (arg) do + env[var] = nil + end + f:kvs_put ("lwj.environ", env) + f:kvs_commit () + end +} + +prog:SubCommand { + name = "getenv", + usage = "VAR [VAR...]", + description = "Export environment to all jobs", + handler = function (self, arg) + local f = assert (require 'flux'.new ()) + local env = f:kvs_get ("lwj.environ") or {} + if #arg == 0 then + for k,v in pairs (env) do + print (k.."="..v) + end + end + for _,k in ipairs (arg) do + local v = env[k] or "" + print (k.."="..v) + end + end +} + + + + -- return keys in dir as a table sorted by number local function sorted_keys (dir) local results = {} diff --git a/t/Makefile.am b/t/Makefile.am index b0bfc9af5dc4..f03e83472613 100644 --- a/t/Makefile.am +++ b/t/Makefile.am @@ -62,6 +62,7 @@ TESTS = \ t1105-proxy.t \ t1999-wreck-rcalc.t \ t2000-wreck.t \ + t2000-wreck-env.t \ t2001-jsc.t \ t2002-pmi.t \ t2003-recurse.t \ @@ -141,6 +142,7 @@ check_SCRIPTS = \ t1105-proxy.t \ t1999-wreck-rcalc.t \ t2000-wreck.t \ + t2000-wreck-env.t \ t2001-jsc.t \ t2002-pmi.t \ t2003-recurse.t \ diff --git a/t/t2000-wreck-env.t b/t/t2000-wreck-env.t new file mode 100755 index 000000000000..ee9407eea6bb --- /dev/null +++ b/t/t2000-wreck-env.t @@ -0,0 +1,51 @@ +#!/bin/sh +# + +test_description='Test basic wreck functionality + +Test basic functionality of wreckrun facility. +' + +. `dirname $0`/sharness.sh +SIZE=${FLUX_TEST_SIZE:-4} +test_under_flux ${SIZE} wreck + +# Return the previous jobid +last_job_id() { + flux wreck last-jobid +} +# Return previous job path in kvs +last_job_path() { + flux wreck last-jobid -p +} + +test_expect_success 'flux-wreck: setenv/getenv works' ' + flux wreck setenv FOO=BAR && + flux wreck getenv FOO +' +test_expect_success 'flux-wreck: unsetenv works' ' + flux wreck unsetenv FOO && + test "$(flux wreck getenv FOO)" = "FOO=" +' +test_expect_success 'flux-wreck: setenv all' ' + flux wreck setenv all && + flux env /usr/bin/env | sort | grep -ve FLUX_URI -e HOSTNAME -e ENVIRONMENT > env.expected && + flux wreck getenv | sort > env.output && + test_cmp env.expected env.output +' +test_expect_success 'wreck: global lwj.environ exported to jobs' ' + flux wreck setenv FOO=bar && + test "$(flux wreckrun -n1 printenv FOO)" = "bar" +' +test_expect_success 'wreck: wreckrun exports environment vars not in global env' ' + BAR=baz flux wreckrun -n1 printenv BAR > printenv.out && + test "$(cat printenv.out)" = "baz" +' +test_expect_success 'wreck: wreckrun --skip-env works' ' + ( export BAR=baz && + test_must_fail flux wreckrun --skip-env -n1 printenv BAR > printenv2.out + ) && + test "$(cat printenv2.out)" = "" +' + +test_done