云原生机器学习平台技术综述（编排调度篇）-来也科技

apiVersion: argoproj.io/v1alpha1kind: Workflowmetadata: generateName: output-parameter-spec: entrypoint: output-parameter templates: # 定义2个component的关系 - name: output-parameter steps:    - - name: generate-parameter template: whalesay - - name: consume-parameter template: print-message arguments: parameters: - name: message # 使用第一个step的输出作为当前step的输入 value: "{{steps.generate-parameter.outputs.parameters.hello-param}}" # 定义第一个component，用于生成字符串 - name: whalesay container: image: docker/whalesay:latest command: [sh, -c] # 打印字符串并写入文件 args: ["sleep 1; echo -n hello world > /tmp/hello_world.txt"] outputs: parameters: # 将文件内作为当前步骤的参数输出 - name: hello-param valueFrom: default: "Foobar" # Default value to use if retrieving valueFrom fails. If not provided workflow will fail instead path: /tmp/hello_world.txt # 定义第二个component，用于打印字符串 - name: print-message inputs: parameters: - name: message container: image: docker/whalesay:latest # 接受一个参数并在控制台打印参数 command: [cowsay] args: ["{{inputs.parameters.message}}"]

def main(_): ps_hosts = FLAGS.ps_hosts.split("，") worker_hosts = FLAGS.worker_hosts.split("，")
 # Create a cluster from the parameter server and worker hosts. cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
 # Create and start a server for the local task. server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)
 if FLAGS.job_name == "ps": server.join() elif FLAGS.job_name == "worker": is_chief = (FLAGS.task_index == 0)  # Assigns ops to the local worker by default. with tf.device(tf.train.replica_device_setter( worker_device="/job:worker/task:%d" % FLAGS.task_index, cluster=cluster)): # Compute

# On ps0.example.com:$ python trainer.py \ --ps_hosts=ps0.example.com:2222,ps1.example.com:2222 \ --worker_hosts=worker0.example.com:2222,worker1.example.com:2222 \ --job_name=ps --task_index=0# On ps1.example.com:$ python trainer.py \ --ps_hosts=ps0.example.com:2222,ps1.example.com:2222 \ --worker_hosts=worker0.example.com:2222,worker1.example.com:2222 \ --job_name=ps --task_index=1# On worker0.example.com:$ python trainer.py \ --ps_hosts=ps0.example.com:2222,ps1.example.com:2222 \ --worker_hosts=worker0.example.com:2222,worker1.example.com:2222 \ --job_name=worker --task_index=0# On worker1.example.com:$ python trainer.py \ --ps_hosts=ps0.example.com:2222,ps1.example.com:2222 \ --worker_hosts=worker0.example.com:2222,worker1.example.com:2222 \ --job_name=worker --task_index=1

import osimport sysimport tempfileimport torchimport torch.distributed as distimport torch.nn as nnimport torch.optim as optimimport torch.multiprocessing as mpfrom torch.nn.parallel import DistributedDataParallel as DDPdef setup(rank, world_size): os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '12355' # initialize the process group dist.init_process_group("gloo", rank=rank, world_size=world_size)def cleanup(): dist.destroy_process_group()

import tensorflow as tfimport horovod.tensorflow.keras as hvd
# Horovod: initialize Horovod.# 初始化 Horovod，启动相关线程和MPI线程hvd.init() 
# Horovod: pin GPU to be used to process local rank (one GPU per process)# 依据 local rank 为不同的进程分配不同的GPUgpus = tf.config.experimental.list_physical_devices('GPU')for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)if gpus: tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
(mnist_images, mnist_labels), _ = \ tf.keras.datasets.mnist.load_data(path='mnist-%d.npz' % hvd.rank())
# 切分数据 dataset = tf.data.Dataset.from_tensor_slices( (tf.cast(mnist_images[..., tf.newaxis] / 255.0, tf.float32), tf.cast(mnist_labels, tf.int64)))dataset = dataset.repeat().shuffle(10000).batch(128)
mnist_model = tf.keras.Sequential([ tf.keras.layers.Conv2D(32, [3, 3], activation='relu'), ...... tf.keras.layers.Dense(10, activation='softmax')])
# Horovod: adjust learning rate based on number of GPUs.# 根据Worker的数量增加学习率的大小scaled_lr = 0.001 * hvd.size()opt = tf.optimizers.Adam(scaled_lr)
# Horovod: add Horovod DistributedOptimizer.# 把常规TensorFlow Optimizer通过Horovod包装起来，进而使用 ring-allreduce 来得到平均梯度opt = hvd.DistributedOptimizer( opt, backward_passes_per_step=1, average_aggregated_gradients=True)
# Horovod: Specify `experimental_run_tf_function=False` to ensure TensorFlow# uses hvd.DistributedOptimizer() to compute gradients.mnist_model.compile(loss=tf.losses.SparseCategoricalCrossentropy(), optimizer=opt, metrics=['accuracy'], experimental_run_tf_function=False)
callbacks = [ # 广播初始化，将模型的参数从第一个设备传向其他设备，以保证初始化模型参数的一致性 hvd.callbacks.BroadcastGlobalVariablesCallback(0), hvd.callbacks.MetricAverageCallback(), hvd.callbacks.LearningRateWarmupCallback(initial_lr=scaled_lr, warmup_epochs=3, verbose=1),]
# Horovod: save checkpoints only on worker 0 to prevent other workers from corrupting them.# 只有设备0需要保存模型参数作为checkpointif hvd.rank() == 0: callbacks.append(tf.keras.callbacks.ModelCheckpoint('./checkpoint-{epoch}.h5'))
# Horovod: write logs on worker 0.verbose = 1 if hvd.rank() == 0 else 0
# Train the model.# Horovod: adjust number of steps based on number of GPUs.mnist_model.fit(dataset, steps_per_epoch=500 // hvd.size(), callbacks=callbacks, epochs=24, verbose=verbose)

# To run on 4 machines with 4 GPUs each:$ horovodrun -np 16 -H server1:4,server2:4,server3:4,server4:4 python train.py

apiVersion: kubeflow.org/v1kind: TFJobmetadata: generateName: tfjob namespace: your-user-namespacespec: tfReplicaSpecs: PS: replicas: 2 restartPolicy: OnFailure template: metadata: annotations: sidecar.istio.io/inject: "false" spec: containers: - name: tensorflow image: gcr.io/your-project/your-image command: - python - -m - trainer.task - --batch_size=32 - --training_steps=1000 Worker: replicas: 2 restartPolicy: OnFailure template: metadata: annotations: sidecar.istio.io/inject: "false" spec: containers: - name: tensorflow image: gcr.io/your-project/your-image command: - python - -m - trainer.task - --batch_size=32 - --training_steps=1000

// TF_CONFIG{ "cluster":{ "ps": [ "tfjob-ps-0.your-user-namespace.svc:2222", "tfjob-ps-1.your-user-namespace.svc:2222" ], "worker":[ "tfjob-worker-0.your-user-namespace.svc:2222", "tfjob-worker-1.your-user-namespace.svc:2222" ] }, "task":{ "type":"ps", "index":0 }}