def test_build_tf_config(): assert training._build_tf_config(HOST_LIST, HOST1) == { 'cluster': CLUSTER_WITH_PS, 'environment': 'cloud', 'task': MASTER_TASK } assert training._build_tf_config(HOST_LIST, HOST1, ps_task=True) == { 'cluster': CLUSTER_WITH_PS, 'environment': 'cloud', 'task': PS_TASK_1 } assert training._build_tf_config(HOST_LIST, HOST2) == { 'cluster': CLUSTER_WITH_PS, 'environment': 'cloud', 'task': WORKER_TASK } assert training._build_tf_config(HOST_LIST, HOST2, ps_task=True) == { 'cluster': CLUSTER_WITH_PS, 'environment': 'cloud', 'task': PS_TASK_2}
def test_build_tf_config(): assert training._build_tf_config(HOST_LIST, HOST1) == { "cluster": CLUSTER_WITH_PS, "environment": "cloud", "task": MASTER_TASK, } assert training._build_tf_config(HOST_LIST, HOST1, ps_task=True) == { "cluster": CLUSTER_WITH_PS, "environment": "cloud", "task": PS_TASK_1, } assert training._build_tf_config(HOST_LIST, HOST2) == { "cluster": CLUSTER_WITH_PS, "environment": "cloud", "task": WORKER_TASK, } assert training._build_tf_config(HOST_LIST, HOST2, ps_task=True) == { "cluster": CLUSTER_WITH_PS, "environment": "cloud", "task": PS_TASK_2, }
def test_build_tf_config_error(): with pytest.raises(ValueError) as error: training._build_tf_config([HOST1], HOST1, ps_task=True) assert 'Cannot have a ps task if there are no parameter servers in the cluster' in str( error.value)