Package: tf-yarn
tf-yarntf-yarn is a Python library we have built at Criteo for training Pytorch and TensorFlow models on a Hadoop/YARN cluster. An introducing blog post can be foundhere.It supports mono and multi-worker training, different distribution strategies and can run on CPUs or GPUs with just a few lines of code.Prerequisitestf-yarn only supports Python ≥3.6.InstallationNote that tf-yarn does not directly depends on the ML frameworks it supports (TensorFlow, torch ...). That way, TensorFlow users don't install torch and conversely by installing tf-yarn. So you must install the ML framework(s) that you use separately (pip install tensorflow,pip install torch...).Install with Pip$pipinstalltf-yarnInstall from source$gitclonehttps://github.com/criteo/tf-yarn
$cdtf-yarn
$pipinstall.TensorFlow prerequisitesSupported versions: [1.15.0 to 2.2.0].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: OKQuick startDistributing the training of a model with tf-yarn can be decomposed in two steps:Describe your experiment: write the code that will be executed by the workers involved in the training. This includes the instantiation of the model to train, the training dataset (optionally the validation dataset) and the training loop.Run your experiment: execute your code on yarn.Refer to the part dedicated to your ML framework (TensorFlow, Pytorch ...) for a detailed description of these two stepsTensorFlowtf-yarn supports Keras API and the Estimator API (which was the only high-level API of the first TensorFlow releases).Describe your experimentKeras APIA Keras experiment is described by an instance oftf_yarn.tensorflow.KerasExperimentcomposed of the following elements:model: compiled Keras model to trainmodel_dir: hdfs directory where the model will be checkpointedtrain_params: training parameters that will be provided tomodel.fit. This does not include the training examples (input and target data)input_data_fn: function returning the input data (features only) to train the model ontarget_data_fn: function returning the target data (labels only) to train the model onvalidation_data_fn: function returning the data to evaluate the model onExample:fromtf_yarn.tensorflowimportKerasExperimentdefinput_data_fn():dataset=...returndataset.shuffle(1000).batch(128).repeat()defvalidation_data_fn():dataset=...returndataset.shuffle(1000).batch(128)model=tf.keras.Sequential()...opt=tf.keras.optimizers.Adadelta(1.0*HVD_SIZE)model.compile(loss='sparse_categorical_crossentropy',optimizer=opt,metrics=['accuracy'])train_params={"steps_per_epoch":100,"callbacks":my_callbacks}experiment=KerasExperiment(model=model,model_dir=hdfs_dir,train_params=train_params,input_data_fn=input_data_fn,target_data_fn=None,validation_data_fn=validation_data_fn)Estimator APIThe experiment is described by an instance oftf_yarn.tensorflow.Experimentcomposed of the following elements:estimator: the model to traintrain_spec: an instance oftf.estimator.TrainSpeceval_spec: an instance oftf.estimator.EvalSpecfromtf_yarn.tensorflowimportExperimentestimator=tf.estimator.Estimator(model_fn=model_fn)train_spec=tf.estimator.TrainSpec(input_fn,max_steps=1)eval_spec=tf.estimator.EvalSpec(input_fn,steps=1)experiment=Experiment(estimator,train_spec,eval_spec)Run your experimentTo run your experiment on yarn, simply call the methodtf_yarn.tensorflow.run_on_yarn. The only mandatory parameter is experiment_fn which must be a function accepting no parameter and returning your objecttf_yarn.tensorflow.KerasExperimentortf_yarn.tensorflow.Experimentwhich describes your experiment.fromtf_yarn.tensorflowimportrun_on_yarn,KerasExperimentdefexperiment_fn():...returnKerasExperiment(model=model,model_dir=hdfs_dir,train_params=train_params,input_data_fn=input_data_fn,target_data_fn=None,validation_data_fn=validation_data_fn)run_on_yarn(experiment_fn)The default distribution strategy isParameterServerStrategywhich belongs to the group of asynchronous distribution strategies.
Although this distribution strategy works very well with the Estimator API, we did not manage to make it work with the Keras API (with TensorFlow <= 2.2). So we advise Keras users to usehorovod gloofor distributing the training. Note that horovod gloo is a synchronous distribution strategy based on all-reduce ops:fromtf_yarn.tensorflowimportrun_on_yarn,KerasExperimentdefexperiment_fn():...returnKerasExperiment(model=model,model_dir=hdfs_dir,train_params=train_params,input_data_fn=input_data_fn,target_data_fn=None,validation_data_fn=validation_data_fn)run_on_yarn(experiment_fn,custom_task_module="tf_yarn.tensorflow.tasks.gloo_allred_task")PytorchDescribe your experimentA Pytorch experiment is described by an instance oftf_yarn.pytorch.PytorchExperimentcomposed of the following elements:model: model to trainmain_fn: Main function run to train the model. This function is executed by all workers involved in the training. It must accept these inputs: model to train, train dataloader, device (cpu:0, cpu:1, cuda:0, cuda:1 ...) allocated to the worker for the training and rank (worker id).Training dataset: training dataset (instance oftorch.utils.data.Dataset,webdataset.WebDatasetorwebdataset.DataPipeline).dataloader_args: parameters (batch size, number of workers, collate function ...) passed to the dataloader used to load and iterate over the training dataset. Instance oftf_yarn.pytorch.DataLoaderArgs.tensorboard_hdfs_dir: HDFS directory where tensorboard results will be written at the end of the trainingddp_args: DistributedDataParallel parameters. Refer toPytorch documentation. Instance oftf_yarn.pytorch.DistributedDataParallelArgsfromtf_yarn.pytorchimportPytorchExperimentdefmain_fn(model:torch.nn.Module,trainloader:torch.utils.data.dataloader.DataLoader,device:str,rank:int):loss_fn=nn.CrossEntropyLoss()optimizer=optim.SGD(model.parameters(),lr=0.001,momentum=0.9)forepochinrange(10):trainloader.sampler.set_epoch(epoch)fori,datainenumerate(trainloader,0):data=data.to(rank)prediction=model(data)transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])trainset=torchvision.datasets.CIFAR10(root='./data',train=True,download=True,transform=transform)experiment=PytorchExperiment(model=model,main_fn=main_fn,train_dataset=trainset,dataloader_args=DataLoaderArgs(batch_size=4,num_workers=2))Run your experimentTo run your experiment on yarn, simply call the methodtf_yarn.pytorch.run_on_yarn. The only mandatory parameters are:experiment_fn: must be a function accepting no parameter and returning your objecttf_yarn.pytorch.PytorchExperimentwhich describes your experiment.task_specs: describe yarn resources that you want to use for your experiment.fromtf_yarn.pytorchimportrun_on_yarn,PytorchExperiment,TaskSpecdefexperiment_fn():...returnPytorchExperiment(model=model,main_fn=main_fn,train_dataset=trainset,dataloader_args=DataLoaderArgs(batch_size=4,num_workers=2))run_on_yarn(experiment_fn,task_specs={"worker":TaskSpec(memory=48*2**10,vcores=48,instances=2,nb_proc_per_worker=2,label=NodeLabel.GPU)})The distribution strategy isDistributedDataParallelwhich belongs to the family of synchronous distribution strategies.run_on_yarnThe method run_on_yarn exposes several parameters that let you configure the yarn job that will be created to train your model on yarn:pyenv_zip_path: Path to an archive of your python environment that will be used by executors to run your experiment. It can be a zipped conda env or a pex archive.
If your python environement is different between CPU and GPU, you can provide a dictionnary fromtf_yarn.topologies.NodeLabelto a python environment. Example:fromtf_yarnimportNodeLabel...run_on_yarn(...,pyenv_zip_path={NodeLabel.CPU:"viewfs://root/path/to/env-cpu",NodeLabel.GPU:"viewfs://root/path/to/env-gpu"})If no archive is provided, tf-yarn will automatically package your active python environment in a pex.task_specs: used to define the resources that you need for your experiment. Dictionary from task names (ps, worker, chief, evaluator, tensorboard ...) totf_yarn.topologies.TaskSpec. Example:fromtf_yarnimportTaskSpec,NodeLabel...run_on_yarn(...,task_specs={"worker":TaskSpec(memory=48*2**10,vcores=48,instances=2,label=NodeLabel.GPU),"tensorboard":TaskSpec(memory=16*2**10,vcores=4,instances=1,label=NodeLabel.CPU)})In this example, we are requesting 2 executors with GPUs, 48 vcores and 48 GBs for workers and 1 executor with 4 vcores and 16 GBs for tensorboard.files: local files or directories to upload on the executors. Dictionary from target location (on executor) to local location (on your local machine). Target locations must be relative to the executor root directory. Note that the executor root directory is appended toPYTHONPATH. Therefore, any listed Python module will be importable.env: environment variables to set on executors. Dictionary from variable name to variable value. Example:run_on_yarn(...,env={"HADOOP_CONF_DIR":"/etc/hadoop/conf.all"})queue: yarn queue to schedule your job in. Example:run_on_yarn(...,queue="ml-gpu")acls: configures the application-level Access Control Lists (ACLs). Optional, defaults to ACLs all access. SeeACLs <https://jcrist.github.io/skein/specification.html#acls>for details.file_systems:skeinthe 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. This would instructskeinto acquire a delegation token for
these namenodes in addition tofs.defaultFS:run_on_yarn(...,file_systems=["hdfs://preprod"])nb_retries: number of times the yarn application is retried in case of failuresname: Name of the yarn applicationModel evaluationThis feature is not supported for Pytorch.Model training and model evaluation can be run independently. To do so, you must
use the 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")ExamplesPlease refer to the various examples available inexamplesOther documentationsMLflowto track experiments.
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-reduce