Expand environment variables in pod specs (#101)

dask · Oct 25, 2018 · 09e894d · 09e894d
1 parent 295f064
commit 09e894d
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/dask_kubernetes/core.py b/dask_kubernetes/core.py
@@ -157,6 +157,7 @@ def __init__(
 
         if not pod_template and dask.config.get('kubernetes.worker-template', None):
             d = dask.config.get('kubernetes.worker-template')
+            d = dask.config.expand_environment_variables(d)
             pod_template = make_pod_from_dict(d)
 
         if not pod_template and dask.config.get('kubernetes.worker-template-path', None):
@@ -165,6 +166,7 @@ def __init__(
             fn = fn.format(**os.environ)
             with open(fn) as f:
                 d = yaml.safe_load(f)
+            d = dask.config.expand_environment_variables(d)
             pod_template = make_pod_from_dict(d)
 
         if not pod_template:
@@ -275,6 +277,7 @@ def from_yaml(cls, yaml_path, **kwargs):
             raise ImportError("PyYaml is required to use yaml functionality, please install it!")
         with open(yaml_path) as f:
             d = yaml.safe_load(f)
+            d = dask.config.expand_environment_variables(d)
             return cls.from_dict(d, **kwargs)
 
     @property

diff --git a/dask_kubernetes/tests/test_core.py b/dask_kubernetes/tests/test_core.py
@@ -204,6 +204,42 @@ def test_pod_from_yaml(image_name, loop, ns):
                 assert all(client.has_what().values())
 
 
+def test_pod_from_yaml_expand_env_vars(image_name, loop, ns):
+    try:
+        os.environ["FOO_IMAGE"] = image_name
+
+        test_yaml = {
+            "kind": "Pod",
+            "metadata": {
+                "labels": {
+                    "app": "dask",
+                    "component": "dask-worker"
+                }
+            },
+            "spec": {
+                "containers": [{
+                    "args": [
+                        "dask-worker",
+                        "$(DASK_SCHEDULER_ADDRESS)",
+                        "--nthreads",
+                        "1"
+                    ],
+                    "image": '${FOO_IMAGE}',
+                    'imagePullPolicy': 'IfNotPresent',
+                    "name": "dask-worker"
+                }]
+            }
+        }
+
+        with tmpfile(extension='yaml') as fn:
+            with open(fn, mode='w') as f:
+                yaml.dump(test_yaml, f)
+            with KubeCluster.from_yaml(f.name, loop=loop, namespace=ns) as cluster:
+                assert cluster.pod_template.spec.containers[0].image == image_name
+    finally:
+        del os.environ['FOO_IMAGE']
+
+
 def test_pod_from_dict(image_name, loop, ns):
     spec = {
         'metadata': {},