Merge pull request #1405 from grondo/wreck-setenv

wreck: better integration of global lwj.environ environment for wreck jobs

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,
 

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 = {}

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 \

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