Merge pull request #27 from Autodesk/environment-variables

Set job environment variables
Autodesk · Mar 9, 2018 · 614c0a0 · 614c0a0
2 parents 25ac40d + 724610e
commit 614c0a0
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 3 deletions.
diff --git a/pyccc/engines/dockerengine.py b/pyccc/engines/dockerengine.py
@@ -86,6 +86,9 @@ def _generate_container_args(self, job):
                               working_dir=job.workingdir,
                               environment={'PYTHONIOENCODING':'utf-8'})
 
+        if job.env:
+            container_args['environment'].update(job.env)
+
         if job.engine_options:
             volumes = []
             binds = []

diff --git a/pyccc/engines/subproc.py b/pyccc/engines/subproc.py
@@ -70,6 +70,8 @@ def submit(self, job):
 
         subenv = os.environ.copy()
         subenv['PYTHONIOENCODING'] = 'utf-8'
+        if job.env:
+            subenv.update(job.env)
         job.subproc = subprocess.Popen(job.command,
                                        shell=True,
                                        cwd=job.workingdir,

diff --git a/pyccc/job.py b/pyccc/job.py
@@ -71,6 +71,7 @@ class Job(object):
         engine_options (dict): additional engine-specific options
         workingdir (str): working directory in the execution environment (i.e., on the local
             system for a subprocess, or inside the container for a docker engine)
+        env (Dict[str,str]): custom environment variables for the Job
     """
     def __init__(self, engine=None,
                  image=None,
@@ -84,14 +85,16 @@ def __init__(self, engine=None,
                  on_status_update=None,
                  when_finished=None,
                  workingdir=None,
-                 engine_options=None):
+                 engine_options=None,
+                 env=None):
 
         self.name = name
         self.engine = engine
         self.image = image
         self.command = if_not_none(command, '')
         self.engine_options = engine_options
         self.workingdir = workingdir
+        self.env = env
 
         self.inputs = inputs
         if self.inputs is not None:  # translate strings into file objects

diff --git a/pyccc/tests/test_job_types.py b/pyccc/tests/test_job_types.py
@@ -448,7 +448,6 @@ def test_abspath_input_files(fixture, request):
 def test_directory_input(fixture, request):
     engine = request.getfuncargvalue(fixture)
 
-    # this is OK with docker but should fail with a subprocess
     job = engine.launch(image='alpine', command='cat data/a data/b',
                         inputs={'data':
                                     pyccc.LocalDirectoryReference(os.path.join(THISDIR, 'data'))})
@@ -461,7 +460,6 @@ def test_directory_input(fixture, request):
 def test_passing_files_between_jobs(fixture, request):
     engine = request.getfuncargvalue(fixture)
 
-    # this is OK with docker but should fail with a subprocess
     job1 = engine.launch(image='alpine', command='echo hello > world')
     job1.wait()
     assert job1.exitcode == 0
@@ -471,3 +469,15 @@ def test_passing_files_between_jobs(fixture, request):
     job2.wait()
     assert job2.exitcode == 0
     assert job2.stdout.strip() == 'hello'
+
+
+@pytest.mark.parametrize('fixture', fixture_types['engine'])
+def test_job_env_vars(fixture, request):
+    engine = request.getfuncargvalue(fixture)
+
+    job = engine.launch(image='alpine',
+                        command='echo ${AA} ${BB}',
+                        env={'AA': 'hello', 'BB':'world'})
+    job.wait()
+    assert job.exitcode == 0
+    assert job.stdout.strip() == 'hello world'