Merge pull request #163 from OSC/copy_environment

Copy environment

lib/ood_core/job/adapters/lsf/helper.rb

@@ -92,6 +92,15 @@ def batch_submit_args(script, after: [], afterok: [], afternotok: [], afterany:
     args += ["-W", (script.wall_time / 60).to_i] unless script.wall_time.nil?
     args += ["-L", script.shell_path.to_s] unless script.shell_path.nil?
 
+    # environment
+    env = script.job_environment || {}
+    # To preserve pre-existing behavior we only act when true or false, when nil we do nothing
+    if script.copy_environment?
+      args += ["-env", (["all"] + env.keys).join(",")]
+    elsif script.copy_environment? == false
+      args += ["-env", (["none"] + env.keys).join(",")]
+    end
+
     # input and output files
     args += ["-i", script.input_path] unless script.input_path.nil?
     args += ["-o", script.output_path] unless script.output_path.nil?
@@ -104,9 +113,6 @@ def batch_submit_args(script, after: [], afterok: [], afternotok: [], afterany:
 
     args += script.native unless script.native.nil?
 
-    # environment
-    env = script.job_environment || {}
-
     {args: args, env: env}
   end
 end

lib/ood_core/job/adapters/pbspro.rb

@@ -261,6 +261,7 @@ def submit(script, after: [], afterok: [], afternotok: [], afterany: [])
           # Set environment variables
           envvars = script.job_environment.to_h
           args += ["-v", envvars.map{|k,v| "#{k}=#{v}"}.join(",")] unless envvars.empty?
+          args += ["-V"] if script.copy_environment?
 
           # If error_path is not specified we join stdout & stderr (as this
           # mimics what the other resource managers do)

lib/ood_core/job/adapters/sge/helper.rb

@@ -20,6 +20,7 @@ def batch_submit_args(script, after: [], afterok: [], afternotok: [], afterany:
     args += ['-h'] if script.submit_as_hold
     args += ['-r', 'yes'] if script.rerunnable
     script.job_environment.each_pair {|k, v| args += ['-v', "#{k.to_s}=#{v.to_s}"]} unless script.job_environment.nil?
+    args += ["-V"] if script.copy_environment?
 
     if script.workdir
       args += ['-wd', script.workdir]

lib/ood_core/job/adapters/slurm.rb

@@ -192,7 +192,7 @@ def delete_job(id)
           # @return [String] the id of the job that was created
           def submit_string(str, args: [], env: {})
             args = args.map(&:to_s) + ["--parsable"]
-            env = {"SBATCH_EXPORT" => "NONE"}.merge env.each_with_object({}) { |(k, v), h| h[k.to_s] = v.to_s }
+            env = env.to_h.each_with_object({}) { |(k, v), h| h[k.to_s] = v.to_s }
             call("sbatch", *args, env: env, stdin: str.to_s).strip.split(";").first
           end
 
@@ -394,7 +394,10 @@ def submit(script, after: [], afterok: [], afternotok: [], afterany: [])
 
           # Set environment variables
           env = script.job_environment || {}
-          args += ["--export", script.job_environment.keys.join(",")] unless script.job_environment.nil? || script.job_environment.empty?
+          unless (script.job_environment.nil? || script.job_environment.empty?)
+            prefix = script.copy_environment? ? "ALL," : "NONE,"  # NONE if false or nil
+            args += ["--export", prefix + script.job_environment.keys.join(",")]
+          end
 
           # Set native options
           args += script.native if script.native

lib/ood_core/job/adapters/torque.rb

@@ -155,6 +155,7 @@ def submit(script, after: [], afterok: [], afternotok: [], afterany: [])
             # Set environment variables
             env = script.job_environment.to_h
             args += ["-v", env.keys.join(",")] unless env.empty?
+            args += ["-V"] if script.copy_environment?
 
             # If error_path is not specified we join stdout & stderr (as this
             # mimics what the other resource managers do)

lib/ood_core/job/script.rb

@@ -104,6 +104,11 @@ class Script
       # @return [Object, nil] native specifications
       attr_reader :native
 
+      # Flag whether the job should contain a copy of its calling environment
+      # @return [Boolean] copy environment
+      attr_reader :copy_environment
+      alias_method :copy_environment?, :copy_environment
+
       # @param content [#to_s] the script content
       # @param args [Array<#to_s>, nil] arguments supplied to script
       # @param submit_as_hold [Boolean, nil] whether job is held after submit
@@ -132,7 +137,8 @@ def initialize(content:, args: nil, submit_as_hold: nil, rerunnable: nil,
                      job_name: nil, shell_path: nil, input_path: nil,
                      output_path: nil, error_path: nil, reservation_id: nil,
                      queue_name: nil, priority: nil, start_time: nil,
-                     wall_time: nil, accounting_id: nil, job_array_request: nil, native: nil, **_)
+                     wall_time: nil, accounting_id: nil, job_array_request: nil,
+                     native: nil, copy_environment: nil, **_)
         @content = content.to_s
 
         @submit_as_hold      = submit_as_hold
@@ -157,6 +163,7 @@ def initialize(content:, args: nil, submit_as_hold: nil, rerunnable: nil,
         @accounting_id      = accounting_id     && accounting_id.to_s
         @job_array_request  = job_array_request && job_array_request.to_s
         @native             = native
+        @copy_environment   = (copy_environment.nil?) ? nil : !! copy_environment
       end
 
       # Convert object to hash
@@ -184,7 +191,8 @@ def to_h
           wall_time:           wall_time,
           accounting_id:       accounting_id,
           job_array_request:   job_array_request,
-          native:              native
+          native:              native,
+          copy_environment:    copy_environment
         }
       end
 

spec/job/adapters/lsf/helper_spec.rb

@@ -186,5 +186,27 @@ def args_for(attrs = {})
       expect(args_for(email: ["efranz@osc.edu", "efranz2@osc.edu"])).to eq({args: ["-u", "efranz@osc.edu,efranz2@osc.edu"], env: {}})
       expect(args_for(email: [])).to eq({args: [], env: {}})
     end
+
+    it "when Script#copy_environment? is true, and Script#job_environment is empty" do
+      expect(args_for(copy_environment: true)).to eq({args: ["-env", "all"], env: {}})
+    end
+
+    it "when Script#copy_environment? is false, and Script#job_environment is empty" do
+      expect(args_for(copy_environment: false)).to eq({args: ["-env", "none"], env: {}})
+    end
+
+    it "when Script#copy_environment? is true, and Script#job_environment is not empty" do
+      env = {"FOO" => 'BAR'}
+      expect(args_for(copy_environment: true, job_environment: env)).to eq({args: ["-env", "all,FOO"], env: env})
+    end
+
+    it "when Script#copy_environment? is false, and Script#job_environment is not empty" do
+      env = {"FOO" => 'BAR'}
+      expect(args_for(copy_environment: false, job_environment: env)).to eq({args: ["-env", "none,FOO"], env: env})
+    end
+
+    it "when Script#copy_environment? is nil" do
+      expect(args_for(copy_environment: nil)).to eq({args: [], env: {}})
+    end
   end
 end

spec/job/adapters/pbspro_spec.rb

@@ -104,6 +104,12 @@ def build_script(opts = {})
       it { expect(pbspro).to have_received(:submit_string).with(content, args: ["-v", "key=value", "-j", "oe"], chdir: nil) }
     end
 
+    context "with :job_environment and Script#copy_environment is true" do
+      before { adapter.submit(build_script(copy_environment: true, job_environment: {"key" => "value"})) }
+
+      it { expect(pbspro).to have_received(:submit_string).with(content, args: ["-v", "key=value", "-V", "-j", "oe"], chdir: nil) }
+    end
+
     context "with :workdir" do
       before { adapter.submit(build_script(workdir: "/path/to/workdir")) }
 

spec/job/adapters/sge_spec.rb

@@ -86,6 +86,12 @@ def build_script(opts = {})
       it { expect(batch).to have_received(:submit).with(content, ["-v", "key=value", "-cwd"]) }
     end
 
+    context "with :job_environment and Script#copy_environment" do
+      before { adapter.submit(build_script(copy_environment: true, job_environment: {"key" => "value"})) }
+
+      it { expect(batch).to have_received(:submit).with(content, ["-v", "key=value", "-V", "-cwd"]) }
+    end
+
     context "with :workdir" do
       before { adapter.submit(build_script(workdir: "/path/to/workdir")) }
 

spec/job/adapters/slurm_spec.rb

@@ -100,7 +100,13 @@ def build_script(opts = {})
     context "with :job_environment" do
       before { adapter.submit(build_script(job_environment: {"key" => "value"})) }
 
-      it { expect(slurm).to have_received(:submit_string).with(content, args: ["--export", "key"], env: {"key" => "value"}) }
+      it { expect(slurm).to have_received(:submit_string).with(content, args: ["--export", "NONE,key"], env: {"key" => "value"}) }
+    end
+
+    context "with :job_environment where copy_environment is true" do
+      before { adapter.submit(build_script(copy_environment: true, job_environment: {"key" => "value"})) }
+
+      it { expect(slurm).to have_received(:submit_string).with(content, args: ["--export", "ALL,key"], env: {"key" => "value"}) }
     end
 
     context "with :workdir" do

spec/job/adapters/torque_spec.rb

@@ -103,6 +103,12 @@ def build_script(opts = {})
       it { expect(pbs).to have_received(:submit).with(content, args: ["-v", "key", "-j", "oe"], env: {"key" => "value"}, chdir: nil) }
     end
 
+    context "with :job_environment and Script#copy_environment = true" do
+      before { adapter.submit(build_script(copy_environment: true, job_environment: {"key" => "value"})) }
+
+      it { expect(pbs).to have_received(:submit).with(content, args: ["-v", "key", "-V", "-j", "oe"], env: {"key" => "value"}, chdir: nil) }
+    end
+
     context "with :workdir" do
       before { adapter.submit(build_script(workdir: "/path/to/workdir")) }
 

spec/job/script_spec.rb

@@ -38,6 +38,7 @@ def build_script(opts = {})
   it { is_expected.to respond_to(:job_array_request) }
   it { is_expected.to respond_to(:native) }
   it { is_expected.to respond_to(:to_h) }
+  it { is_expected.to respond_to(:copy_environment) }
 
   describe '.new' do
     context "when :context not defined" do
@@ -209,6 +210,7 @@ def build_script(opts = {})
     it { is_expected.to have_key(:accounting_id) }
     it { is_expected.to have_key(:job_array_request) }
     it { is_expected.to have_key(:native) }
+    it { is_expected.to have_key(:copy_environment) }
   end
 
   describe "#==" do
@@ -220,4 +222,18 @@ def build_script(opts = {})
       is_expected.not_to eq(build_script(priority: 123))
     end
   end
+
+  describe "#copy_environment" do
+    context "when default constructed" do
+      subject { build_script().copy_environment }
+
+      it { is_expected.to eq(nil) }
+    end
+
+    context "when set to be true" do
+      subject { build_script(copy_environment: true).copy_environment }
+
+      it { is_expected.to eq(true) }
+    end
+  end
 end