Package: tf-yarn-gpu
tf-yarnᵝtf-yarn is a Python library we have built at Criteo for training TensorFlow models on a Hadoop/YARN cluster. An introducing blog post can be foundhere.It supports running on one worker or on multiple workers with different distribution strategies and it can run on CPUs or GPUs using just a few lines of code.Its API provides an easy entry point for working with Estimators and Keras. Please refer to theexamplesfor some code samples.MLflowis supported for all kind of trainings (one worker/distributed).
More infoshere.Tensorboardcan be spawned in a separate container during learnings.Two alternatives to TensorFlow's distribution strategies are available:Horovod with glooandtf-collective-all-reduceInstallationInstall with PipNote that in order to support both tensorflow cpu and gpu in a single package, tf-yarn doesn't depend directly on tensorflow.
You can either runpip install tf-yarnthen install tensorflow or tensorflow_gpu, or usepip install tf-yarn[cpu]to do it all in one command.
The supported versions of tensorflow are in the range 1.15.0 to 2.2.0. Install a version in that range to use tf yarn.$pipinstalltf-yarn[cpu]Install from source$gitclonehttps://github.com/criteo/tf-yarn
$cdtf-yarn
$pipinstall.Prerequisitestf-yarn only supports Python ≥3.6.Make sure to have Tensorflow working with HDFS by setting up all the environment variables as describedhere.You can run thecheck_hadoop_envscript to check that your setup is OK (it has been installed by tf_yarn):$ check_hadoop_env
# You should see something like
# INFO:tf_yarn.bin.check_hadoop_env:results will be written in /home/.../shared/Dev/tf-yarn/check_hadoop_env.log
# INFO:tf_yarn.bin.check_hadoop_env:check_env: True
# INFO:tf_yarn.bin.check_hadoop_env:write dummy file to hdfs hdfs://root/tmp/a1df7b99-fa47-4a86-b5f3-9bc09019190f/hello_tf_yarn.txt
# INFO:tf_yarn.bin.check_hadoop_env:check_local_hadoop_tensorflow: True
# INFO:root:Launching remote check
# ...
# INFO:tf_yarn.bin.check_hadoop_env:remote_check: True
# INFO:tf_yarn.bin.check_hadoop_env:Hadoop setup: OKrun_on_yarnThe only abstraction tf-yarn adds on top of the ones already present in
TensorFlow isexperiment_fn. It is a function returning a triple (wrapped in anExperimentobject) of oneEstimatorand two specs --TrainSpecandEvalSpec.Here is a stripped downexperiment_fnfrom one of the providedexamplesto give you an idea of how it might look:fromtf_yarnimportExperimentdefexperiment_fn():# ...estimator=tf.estimator.DNNClassifier(...)returnExperiment(estimator,tf.estimator.TrainSpec(train_input_fn,max_steps=...),tf.estimator.EvalSpec(eval_input_fn))An experiment can be scheduled on YARN using the run_on_yarn function which takes three required arguments:pyenv_zip_pathwhich contains the tf-yarn modules and dependencies like TensorFlow to be shipped to the cluster. pyenv_zip_path can be generated easily with a helper function based on the current installed virtual environment;experiment_fnas described above;task_specsdictionary specifying how much resources to allocate for each of the distributed TensorFlow task type.The example uses theWine Qualitydataset from UCI ML repository. With just under 5000 training instances available, there is no need for multi-node training, meaning that a chief complemented by an evaluator would manage just fine. Note that each task will be executed in its own YARN container.fromtf_yarnimportTaskSpec,run_on_yarnimportcluster_packpyenv_zip_path,_=cluster_pack.upload_env()run_on_yarn(pyenv_zip_path,experiment_fn,task_specs={"chief":TaskSpec(memory="2 GiB",vcores=4),"evaluator":TaskSpec(memory="2 GiB",vcores=1),"tensorboard":TaskSpec(memory="2 GiB",vcores=1)})The final bit is to forward thewinequality.pymodule to the YARN containers,
in order for the tasks to be able to import them:run_on_yarn(...,files={os.path.basename(winequality.__file__):winequality.__file__,})Under the hood, the experiment function is shipped to each container, evaluated and then passed to thetrain_and_evaluatefunction.experiment=experiment_fn()tf.estimator.train_and_evaluate(experiment.estimator,experiment.train_spec,experiment.eval_spec)Specificities using native Keras models instead of estimatorsWhen using a Keras model that is not converted into an estimator,experiment_fnreturns a tuple (wrapped in aKerasExperimentobject) composed of the following elements:model: the compiled Keras modelmodel_dir: the location at which the model and its checkpoints will be savedtrain_params: parameters that will be passed to the model fit method exluding input and target datainput_data_fn: function returning input data for the model fit methodtarget_data_fn: function returning target data for the model fit methodvalidation_data_fn: function returning input data for the model evaluate methodCurrently, Keras models are only supported using Horovod with Gloo as a distribution strategy (and not using MultiWorkerMirroredStrategy). Moreover, Keras models are only supported using Tensorflow 2. We provide an example describing how to use a Keras model with Horovod [examples][native_keras_with_gloo_example].Distributed TensorFlowThe following is a brief summary of the core distributed TensorFlow concepts relevant to trainingestimatorswith the ParameterServerStrategy, as it is the distribution strategy activated by default when training Estimators on multiple nodes.Distributed TensorFlow operates in terms of tasks.
A task has a type which defines its purpose in the distributed TensorFlow cluster:workertasks headed by thechiefdoing model trainingchieftask additionally handling checkpoints, saving/restoring the model, etc.pstasks (aka parameter servers) storing the model itself. These tasks typically do not compute anything.
Their sole purpose is serving the model variablesevaluatortask periodically evaluating the model from the saved checkpointThe types of tasks can depend on the distribution strategy, for example, ps tasks are only used by ParameterServerStrategy.
The following picture presents an example of a cluster setup with 2 workers, 1 chief, 1 ps and 1 evaluator.+-----------+              +---------+   +----------+   +----------+
| evaluator |        +-----+ chief:0 |   | worker:0 |   | worker:1 |
+-----+-----+        |     +----^----+   +-----^----+   +-----^----+
      ^              |          |            |              |
      |              v          |            |              |
      |        +-----+---+      |            |              |
      |        | model   |   +--v---+        |              |
      +--------+ exports |   | ps:0 <--------+--------------+
               +---------+   +------+The cluster is defined by a ClusterSpec, a mapping from task types to their associated network addresses. For instance, for the above example, it looks like that:{
  "chief": ["chief.example.com:2125"],
  "worker": ["worker0.example.com:6784",
             "worker1.example.com:6475"],
  "ps": ["ps0.example.com:7419"],
  "evaluator": ["evaluator.example.com:8347"]
}Starting a task in the cluster requires a ClusterSpec. This means that the spec should be fully known before starting any of the tasks.Once the cluster is known, we need to export the ClusterSpec through theTF_CONFIGenvironment variable and start the TensorFlow server on each container.Then we can run thetrain-and-evaluatefunction on each container.
We just launch the same function as in local training mode, TensorFlow will automatically detect that we have set up a ClusterSpec and start a distributed learning.You can find more information about distributed Tensorflowhereand about distributed training Estimatorshere.Training with multiple workersActivating the previous example in tf-yarn is just changing the cluster_spec by adding the additionalworkerandpsinstances:run_on_yarn(...,task_specs={"chief":TaskSpec(memory="2 GiB",vcores=4),"worker":TaskSpec(memory="2 GiB",vcores=4,instances=2),"ps":TaskSpec(memory="2 GiB",vcores=8),"evaluator":TaskSpec(memory="2 GiB",vcores=1),"tensorboard":TaskSpec(memory="2 GiB",vcores=1)})Configuring the Python interpreter and packagestf-yarn usescluster-packto to ship an isolated virtual environment to the containers.
(You should have installed the dependencies fromrequirements.txtinto your virtual environment firstpip install -r requirements.txt)
This works if you use Anaconda and also withVirtual Environments.By default the generated package is apexpackage. cluster-pack will generate the pex package, upload it to hdfs and you can start tf_yarn by providing the hdfs path.importcluster_packpyenv_zip_path,env_name=cluster_pack.upload_env()run_on_yarn(pyenv_zip_path=pyenv_zip_path)If you hosting evironment is Anacondaupload_envthe packaging module will useconda-packto create the package.You can also directly use the command line tools provided byconda-packandpexto generate the packages.For pex you can run this command in the root directory to create the package (it includes all requirements from setup.py)pex . -o myarchive.pexYou can then run tf-yarn with your generated package:run_on_yarn(pyenv_zip_path="myarchive.pex")Running on GPUYARN does not have first-class support for GPU resources. A common workaround is
to usenode labelswhere CPU-only nodes are unlabelled, while
the GPU ones have a label. Furthermore, in this setting GPU nodes are
typically bound to a separate queue which is different from the default one.Currently, tf-yarn assumes that the GPU label is"gpu". There are no
assumptions on the name of the queue with GPU nodes, however, for the sake of
example we wil use the name"ml-gpu".The default behaviour ofrun_on_yarnis to run on CPU-only nodes. In order
to run on the GPU ones:Set thequeueargument.SetTaskSpec.labeltoNodeLabel.GPUfor relevant task types.
A good rule of a thumb is to run compute heavy"chief"and"worker"tasks on GPU, while keeping"ps"and"evaluator"on CPU.importgetpassimportcluster_packfromtf_yarnimportNodeLabelpyenv_zip_path,_=cluster_pack.upload_env()run_on_yarn(pyenv_zip_pathexperiment_fn,task_specs={"chief":TaskSpec(memory="2 GiB",vcores=4,label=NodeLabel.GPU),"evaluator":TaskSpec(memory="1 GiB",vcores=1)},queue="ml-gpu")The previous example applies to TensorFlow >= 1.15.
For TensorFlow < 1.15 you need to call upload_env with tensorflow-gpu package and provide it torun_on_yarn.Accessing HDFS in the presence offederationskeinthe library underlyingtf_yarnautomatically acquires a delegation token
forfs.defaultFSon security-enabled clusters. This should be enough for most
use-cases. However, if your experiment needs to access data on namenodes other than
the default one, you have to explicitly list them in thefile_systemsargument
torun_on_yarn. This would instructskeinto acquire a delegation token for
these namenodes in addition tofs.defaultFS:run_on_yarn(...,file_systems=["hdfs://preprod"])Depending on the cluster configuration, you might need to point libhdfs to a
different configuration folder. For instance:run_on_yarn(...,env={"HADOOP_CONF_DIR":"/etc/hadoop/conf.all"})Running model evaluation independentlyModel training and model evaluation can be run independently. To do so, you must
use parametercustom_task_moduleofrun_on_yarn.To run model training without evaluation:run_on_yarn(...,task_specs={"chief":TaskSpec(memory="2 GiB",vcores=4),"worker":TaskSpec(memory="2 GiB",vcores=4,instances=2),"ps":TaskSpec(memory="2 GiB",vcores=8),"tensorboard":TaskSpec(memory="2 GiB",vcores=1)})To run model evaluation:run_on_yarn(...,task_specs={"evaluator":TaskSpec(memory="2 GiB",vcores=1)},custom_task_module="tf_yarn.tasks.evaluator_task")