Repository: facebookresearch/hydra

Issue:
Allow setting environment variable automatically for jobs
This is low priority at the moment, please vote up if you will find it useful.

Control of env variables for launched jobs.
This could be based on a special node in the job config:
```
env:
  OMP_NUM_THREADS: 1 
  # copy env from running machine
  USER: ${env:USER} 
```

TO BE CLEAR:
You can already use environment variables in configs:
See the [OmegaConf documentation](https://omegaconf.readthedocs.io/en/1.4_branch/usage.html#environment-variable-interpolation).

This issue is about allowing easier control of the environment variable of Hydra apps for other purposes (for example, setting OMP_NUM_THREADS automatically to 1).


Discussion:
Will this also allow overriding config from environment variable? similar  to overriding from the command line  
This is already possible (although not similar to overriding from the command line).
See the [OmegaConf documentation](https://omegaconf.readthedocs.io/en/1.4_branch/usage.html#environment-variable-interpolation).
I think it will be a good idea to make overriding config from the environment variable consistent with overriding using the command line approach.
How can you make two things that are completely different consistent with each other?
If you want me to take this suggestion seriously, make it more concrete. I have no idea what kind of functionality you are thinking about.

regardless, please open a new issue against OmegaConf if you have a feature request. this is not the right place to discuss any such requests.
OMG, this would be a super important feature. It's hard to believe it doesn't exist already...
Hi, I'm sorry. it's an important enhancement. if we looked at containers, the main workflow would be using environment variables. 
It would be great to have this enhancement soon.
1. This is a feature request to support automatically SETTING the environment variables for Hydra jobs.
2. You can already USE environment variables in the config (see my edit to the task description).
3. I am planning to add this functionality for the next Hydra version.
4. This is marked as "good first issue", meaning it's not hard to implement.

If this is so important to people, they can step up and help.
>  the environment variables for Hydra jobs.

Oh, miss that. 

> You can already USE environment variables in the config

Thank you
It was not clear from the issue description, I updated it.
I recommend everyone here spend the time to read the OmegaConf documentation. there are many good features there you should know about if you use Hydra.

Resolving patch:
diff --git a/hydra/_internal/config_loader_impl.py b/hydra/_internal/config_loader_impl.py
--- a/hydra/_internal/config_loader_impl.py
+++ b/hydra/_internal/config_loader_impl.py
@@ -3,6 +3,7 @@
 Configuration loader
 """
 import copy
+import os
 from typing import List, Optional, Tuple
 
 from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict
@@ -132,6 +133,9 @@ def is_hydra(x: str) -> bool:
             )
             cfg.hydra.job.config_name = config_name
 
+            for key in cfg.hydra.job.env_copy:
+                cfg.hydra.job.env_set[key] = os.environ[key]
+
         return cfg
 
     def load_sweep_config(
diff --git a/hydra/conf/__init__.py b/hydra/conf/__init__.py
--- a/hydra/conf/__init__.py
+++ b/hydra/conf/__init__.py
@@ -1,6 +1,6 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 from dataclasses import dataclass, field
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 from omegaconf import MISSING
 
@@ -58,35 +58,41 @@ class OverridesConf:
     task: List[str] = field(default_factory=lambda: [])
 
 
-@dataclass
 # job runtime information will be populated here
+@dataclass
 class JobConf:
-    # Job name, can be specified by the user (in config or cli) or populated automatically
+    # Job name, populated automatically unless specified by the user (in config or cli)
     name: str = MISSING
 
+    # Populated automatically by Hydra.
     # Concatenation of job overrides that can be used as a part
     # of the directory name.
-    # This can be configured in hydra.job.config.override_dirname
+    # This can be configured via hydra.job.config.override_dirname
     override_dirname: str = MISSING
 
     # Job ID in underlying scheduling system
     id: str = MISSING
 
     # Job number if job is a part of a sweep
-    num: str = MISSING
+    num: int = MISSING
 
     # The config name used by the job
     config_name: Optional[str] = MISSING
 
-    @dataclass
+    # Environment variables to set remotely
+    env_set: Dict[str, str] = field(default_factory=dict)
+    # Environment variables to copy from the launching machine
+    env_copy: List[str] = field(default_factory=list)
+
     # Job config
+    @dataclass
     class JobConfig:
         @dataclass
         # configuration for the ${hydra.job.override_dirname} runtime variable
         class OverrideDirname:
             kv_sep: str = "="
             item_sep: str = ","
-            exclude_keys: List[str] = field(default_factory=lambda: [])
+            exclude_keys: List[str] = field(default_factory=list)
 
         override_dirname: OverrideDirname = OverrideDirname()
 
diff --git a/hydra/core/utils.py b/hydra/core/utils.py
--- a/hydra/core/utils.py
+++ b/hydra/core/utils.py
@@ -4,6 +4,7 @@
 import os
 import re
 import sys
+from contextlib import contextmanager
 from dataclasses import dataclass
 from os.path import basename, dirname, splitext
 from pathlib import Path
@@ -130,7 +131,8 @@ def run_job(
         _save_config(task_cfg, "config.yaml", hydra_output)
         _save_config(hydra_cfg, "hydra.yaml", hydra_output)
         _save_config(config.hydra.overrides.task, "overrides.yaml", hydra_output)
-        ret.return_value = task_function(task_cfg)
+        with env_override(hydra_cfg.hydra.job.env_set):
+            ret.return_value = task_function(task_cfg)
         ret.task_name = JobRuntime.instance().get("name")
 
         # shut down logging to ensure job log files are closed.
@@ -212,3 +214,20 @@ def split_config_path(
         config_name = config_file
 
     return config_dir, config_name
+
+
+@contextmanager
+def env_override(env: Dict[str, str]) -> Any:
+    """Temporarily set environment variables inside the context manager and
+    fully restore previous environment afterwards
+    """
+    original_env = {key: os.getenv(key) for key in env}
+    os.environ.update(env)
+    try:
+        yield
+    finally:
+        for key, value in original_env.items():
+            if value is None:
+                del os.environ[key]
+            else:
+                os.environ[key] = value


Test patch:
diff --git a/hydra/test_utils/test_utils.py b/hydra/test_utils/test_utils.py
--- a/hydra/test_utils/test_utils.py
+++ b/hydra/test_utils/test_utils.py
@@ -255,27 +255,45 @@ def verify_dir_outputs(
     ) == OmegaConf.create(overrides or [])
 
 
+def _get_statements(indent: str, statements: Union[None, str, List[str]]) -> str:
+    if isinstance(statements, str):
+        statements = [statements]
+
+    code = ""
+    if statements is None or len(statements) == 0:
+        code = "pass"
+    else:
+        for p in statements:
+            code += f"{indent}{p}\n"
+    return code
+
+
 def integration_test(
     tmpdir: Path,
     task_config: DictConfig,
     overrides: List[str],
     prints: Union[str, List[str]],
     expected_outputs: Union[str, List[str]],
+    prolog: Union[None, str, List[str]] = None,
     filename: str = "task.py",
     env_override: Dict[str, str] = {},
+    clean_environment: bool = False,
 ) -> str:
-    if isinstance(prints, str):
-        prints = [prints]
     if isinstance(expected_outputs, str):
         expected_outputs = [expected_outputs]
     if isinstance(task_config, (list, dict)):
         task_config = OmegaConf.create(task_config)
+    if isinstance(prints, str):
+        prints = [prints]
+    prints = [f'f.write({p} + "\\n")' for p in prints]
 
     s = string.Template(
         """import hydra
 import os
 from hydra.core.hydra_config import HydraConfig
 
+$PROLOG
+
 @hydra.main($CONFIG_PATH)
 def experiment(cfg):
     with open("$OUTPUT_FILE", "w") as f:
@@ -286,14 +304,8 @@ def experiment(cfg):
 """
     )
 
-    print_code = ""
-    print_indent = "        "
-    if prints is None or len(prints) == 0:
-        print_code = "pass"
-    else:
-        for p in prints:
-            print_code += f"{print_indent}f.write({p})\n"
-            print_code += f'{print_indent}f.write("\\n")\n'
+    print_code = _get_statements(indent="        ", statements=prints)
+    prolog_code = _get_statements(indent="", statements=prolog)
 
     config_path = ""
     if task_config is not None:
@@ -304,7 +316,10 @@ def experiment(cfg):
     # replace Windows path separator \ with an escaped version \\
     output_file = output_file.replace("\\", "\\\\")
     code = s.substitute(
-        PRINTS=print_code, CONFIG_PATH=config_path, OUTPUT_FILE=output_file
+        PRINTS=print_code,
+        CONFIG_PATH=config_path,
+        OUTPUT_FILE=output_file,
+        PROLOG=prolog_code,
     )
     task_file = tmpdir / filename
     task_file.write_text(str(code), encoding="utf-8")
@@ -313,8 +328,11 @@ def experiment(cfg):
     orig_dir = os.getcwd()
     try:
         os.chdir(str(tmpdir))
-        modified_env = os.environ.copy()
-        modified_env.update(env_override)
+        if clean_environment:
+            modified_env = {}
+        else:
+            modified_env = os.environ.copy()
+            modified_env.update(env_override)
         subprocess.check_call(cmd, env=modified_env)
 
         with open(output_file, "r") as f:
@@ -324,7 +342,7 @@ def experiment(cfg):
         if expected_outputs is not None:
             assert len(output) == len(
                 expected_outputs
-            ), f"Unexpected number of output lines from {task_file}"
+            ), f"Unexpected number of output lines from {task_file}, output lines:\n\n{file_str}"
             for idx in range(len(output)):
                 assert (
                     output[idx] == expected_outputs[idx]
diff --git a/tests/test_apps/setting_env/__init__.py b/tests/test_apps/setting_env/__init__.py
new file mode 100644
--- /dev/null
+++ b/tests/test_apps/setting_env/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
diff --git a/tests/test_apps/setting_env/config.yaml b/tests/test_apps/setting_env/config.yaml
new file mode 100644
--- /dev/null
+++ b/tests/test_apps/setting_env/config.yaml
@@ -0,0 +1,9 @@
+dataset:
+  name: imagenet
+  path: /datasets/imagenet
+
+hydra:
+  job:
+    env:
+      foo: bar
+      bar: ${hydra.job.num}
\ No newline at end of file
diff --git a/tests/test_apps/setting_env/my_app.py b/tests/test_apps/setting_env/my_app.py
new file mode 100644
--- /dev/null
+++ b/tests/test_apps/setting_env/my_app.py
@@ -0,0 +1,16 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import os
+
+from omegaconf import DictConfig
+
+import hydra
+
+
+@hydra.main(config_path="config.yaml")
+def my_app(_: DictConfig) -> None:
+    print(f"foo={os.environ['foo']}")
+    print(f"bar={os.environ['bar']}")
+
+
+if __name__ == "__main__":
+    my_app()
diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py
--- a/tests/test_config_loader.py
+++ b/tests/test_config_loader.py
@@ -11,6 +11,7 @@
 from hydra.core.config_loader import LoadTrace
 from hydra.core.config_store import ConfigStore, ConfigStoreWithProvider
 from hydra.core.errors import HydraException
+from hydra.core.utils import env_override
 from hydra.errors import MissingConfigException
 from hydra.test_utils.test_utils import chdir_hydra_root
 
@@ -587,3 +588,12 @@ def test_invalid_plugin_merge(restore_singletons: Any) -> Any:
     cl = ConfigLoaderImpl(config_search_path=create_config_search_path(None))
     with pytest.raises(HydraException):
         cl.load_configuration(config_name="config", overrides=["plugin=invalid"])
+
+
+def test_job_env_copy() -> None:
+    config_loader = ConfigLoaderImpl(config_search_path=create_config_search_path(None))
+    with env_override({"zonk": "123456"}):
+        cfg = config_loader.load_configuration(
+            config_name=None, overrides=["hydra.job.env_copy=[zonk]"]
+        )
+        assert cfg.hydra.job.env_set == {"zonk": "123456"}
diff --git a/tests/test_hydra.py b/tests/test_hydra.py
--- a/tests/test_hydra.py
+++ b/tests/test_hydra.py
@@ -6,7 +6,7 @@
 from typing import Any, List
 
 import pytest
-from omegaconf import OmegaConf
+from omegaconf import DictConfig, OmegaConf
 
 from hydra import MissingConfigException
 from hydra.test_utils.test_utils import (
@@ -592,3 +592,45 @@ def test_sys_exit(tmpdir: Path) -> None:
         "hydra.run.dir=" + str(tmpdir),
     ]
     assert subprocess.run(cmd).returncode == 42
+
+
+@pytest.mark.parametrize(  # type: ignore
+    "task_config, overrides, expected_dir",
+    [
+        ({"hydra": {"run": {"dir": "foo"}}}, [], "foo"),
+        ({}, ["hydra.run.dir=bar"], "bar"),
+        ({"hydra": {"run": {"dir": "foo"}}}, ["hydra.run.dir=boom"], "boom"),
+        (
+            {
+                "hydra": {"run": {"dir": "foo-${hydra.job.override_dirname}"}},
+                "app": {"a": 1, "b": 2},
+            },
+            ["app.a=20"],
+            "foo-app.a=20",
+        ),
+    ],
+)
+def test_local_run_workdir(
+    tmpdir: Path, task_config: DictConfig, overrides: List[str], expected_dir: str
+) -> None:
+    cfg = OmegaConf.create(task_config)
+    assert isinstance(cfg, DictConfig)
+    expected_dir1 = tmpdir / expected_dir
+    integration_test(
+        tmpdir=tmpdir,
+        task_config=cfg,
+        overrides=overrides,
+        prints="os.getcwd()",
+        expected_outputs=str(expected_dir1),
+    )
+
+
+def test_hydra_env_set(tmpdir: Path) -> None:
+    cfg = OmegaConf.create({"hydra": {"job": {"env_set": {"foo": "bar"}}}})
+    integration_test(
+        tmpdir=tmpdir,
+        task_config=cfg,
+        overrides=[],
+        prints="os.environ['foo']",
+        expected_outputs="bar",
+    )
diff --git a/tests/test_integration.py b/tests/test_integration.py
deleted file mode 100644
--- a/tests/test_integration.py
+++ /dev/null
@@ -1,39 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-from pathlib import Path
-from typing import List
-
-import pytest
-from omegaconf import DictConfig, OmegaConf
-
-from hydra.test_utils.test_utils import integration_test
-
-
-@pytest.mark.parametrize(  # type: ignore
-    "task_config, overrides, expected_dir",
-    [
-        ({"hydra": {"run": {"dir": "foo"}}}, [], "foo"),
-        ({}, ["hydra.run.dir=bar"], "bar"),
-        ({"hydra": {"run": {"dir": "foo"}}}, ["hydra.run.dir=boom"], "boom"),
-        (
-            {
-                "hydra": {"run": {"dir": "foo-${hydra.job.override_dirname}"}},
-                "app": {"a": 1, "b": 2},
-            },
-            ["app.a=20"],
-            "foo-app.a=20",
-        ),
-    ],
-)
-def test_local_run_workdir(
-    tmpdir: Path, task_config: DictConfig, overrides: List[str], expected_dir: str
-) -> None:
-    cfg = OmegaConf.create(task_config)
-    assert isinstance(cfg, DictConfig)
-    expected_dir1 = tmpdir / expected_dir
-    integration_test(
-        tmpdir=tmpdir,
-        task_config=cfg,
-        overrides=overrides,
-        prints="os.getcwd()",
-        expected_outputs=str(expected_dir1),
-    )