def setup_tf_config(cluster_spec):
# Note that "evaluator" does not need a cluster, and "ps" (!)
# surprisingly does not follow the same code path as the rest
# and spawns a server regardless of the "environment" value.
task_type, task_id = get_task_description()
_internal.xset_environ(TF_CONFIG=json.dumps(
{
"cluster": cluster_spec,
"environment": "google" if _is_fake_google_env(task_type) else "",
"task": {
"type": task_type,
"index": task_id
},
}))
def test_xset_environ_failure(monkeypatch):
monkeypatch.setattr(os, "environ", {"foo": "bar"})
with pytest.raises(RuntimeError):
xset_environ(foo="boo")
assert os.environ["foo"] == "bar"
def test_xset_environ(monkeypatch):
monkeypatch.setattr(os, "environ", {})
xset_environ(foo="boo")
assert os.environ["foo"] == "boo"