Provide TF_CONFIG environment variable for distributed TensorFlow

[damienpontifex]

The TensorFlow ClusterConfig can parse worker and parameter server settings from a TF_CONFIG environment variable (see https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/learn/python/learn/estimators/run_config.py#L64-L156)

I was trying to pass it via an environment variable in the job configuration file like so:

"{ 'cluster': { 'ps': $AZ_BATCHAI_PS_HOSTS, 'worker': $AZ_BATCHAI_WORKER_HOSTS }, 'task': { 'index': $AZ_BATCHAI_TASK_INDEX, 'type': '' } }"

Which is kind of fine, but falls down for a few cases:

1. When there are no parameter servers (i.e. single node) the ps hosts should be an empty array, but in this case it's just an empty string.
2. The variables for hosts and workers are comma separated and the TF code parses it as JSON, so would ideally be an array type inside this string.
3. The 'task.type' property can be 'master', 'worker' or 'ps' but that doesn't seem to have an appropriate environment variable and I had to pass the option via command line args

More generally though, providing this configuration via a TF_CONFIG environment variable would significantly lower the bar to get distributed training working in TensorFlow and Azure Batch. It would also simplify command line arg parameters and mean just the appropriate data directories would need to be passed and mean the same arguments could be used across master, worker and ps potentially simplify the tensorflowSettings property further.

[AlexanderYukhanov]

Hi Damien,
Thank you for the feedback! We will figure out how to make use of TF_CONFIG for tensorflow framework. At first glance, we can just introduce dedicated environment variables for using with TF_CONFIG.

[damienpontifex]

Great, thanks for the response. It is a JSON serialised dictionary in an environment variable, but would mean distributed training would ‘just work ™️’.

[damienpontifex]

Looking at Azure Batch AI environment variables it seems this is now available.

[AlexanderYukhanov]

Sorry, the functionality is not released yet.

[wtam]

May I ask if anyone or @damienpontifex know what is the env variable for the master host? I encounter "valueerror: if "cluster" is set in tf_config, it must have one "chief" node."
you can reference here the chief node in the cluster spec
https://www.tensorflow.org/api_docs/python/tf/estimator/RunConfig

[damienpontifex]

@wtam my understanding is you must have the task type and index set appropriately for the chief. In the page you linked this is {'cluster': cluster, 'task': {'type': 'chief', 'index': 0}}) where the cluster variable has three keys: chief, ps and worker.

Without seeing your actual code, seems the minimum requirement is cluster to have {'chief': ['host0:2222']}. You can have a look at the logic in RunConfig to see if there's a case with your setup that you have configured wrong.

[wtam]

@damienpontifex Thanks so much for the respond. Since BatchAI only have these env var on BatchAI, $AZ_BATCHAI_PS_HOSTS, $AZ_BATCHAI_WORKER_HOSTS & $AZ_BATCHAI_TASK_INDEX. I overcome the chief node define issueabove by manually reserve the 1st workerhost as chief node and put into the cluster spec. Now I move a bit forward but encounter another issue from RunConfig below:
ValueError: worker is not a valid task_type in the cluster_spec???????:
<tensorflow.python.training.server_lib.ClusterSpec object at 0x7fdd35049750>
Not sure where goes wrong? My cluster is 3 nodes, 1 node reserved for PS and chief node and another 2 nodes for worker node. Appreciate any comment or suggestion to help me out.

This is the cluster spec for the failed worker
'{"cluster": {"chief": ["10.0.0.4:2223"], "worker_hosts": ["10.0.0.5:2222", "10.0.0.6:2222"], "ps_hosts": ["10.0.0.4:2222"]}, "task": {"index": "1", "type": "worker"}}')

[wtam]

Stupid mistake I made on the cluster spec naming, RunConfig is trying to find worker from my worker_hosts and that why I got the ValueError. For people play around with Estimator Distributed GPU on BatchAI, better wait for its support as the away I reserved the worker node also require me to deduct the $AZ_BATCHAI_TASK_INDEX manually in the cluster spec for the workers.

[llidev]

Hi @damienpontifex, maybe you have already known that, Batch AI is now automatically generating TF_CONFIG env var when running tensorflow job. Would you please try it out and please let us know if it works for you? Thanks!

[yangsiyu007]

Hi @lliimsft, I'm seeing the automatically generated TF_CONFIG env var with nodeCount 1 as:
{'task': {'type': 'master', 'index': 0}, 'cluster': {'ps': [''], 'worker': ['10.0.0.4:2222']}, 'environment': 'cloud'}
which doesn't seem to work in this 1-node cluster scenario?

[damienpontifex]

Getting this error when running with nodeCount=3 in the stderr-ps-0.txt log

"ValueError: If "cluster" is set in TF_CONFIG, it must have one "chief" node."

For this task the TF_CONFIG variable was:

{'cluster': {'worker': ['10.0.0.4:2223', '10.0.0.5:2222', '10.0.0.6:2222'], 'ps': ['10.0.0.4:2222']}, 'task': {'type': 'ps', 'index': 0}, 'environment': 'cloud'}

The worker logs just had "Warning: Permanently added '[10.0.0.5]:23' (ECDSA) to the list of known hosts."

[damienpontifex]

I put the code I'm running here https://github.com/damienpontifex/BatchAIMnist

From the repo, I do:

sh prepare-cluster.sh
sh data-prep.sh
# Wait until data prep done
sh train.sh