text: # Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import unittest
  import subprocess
  import sys, os
  import json
  import shutil
  
  import random
  
  from os import listdir
  from os.path import isfile, join
  
  pyname = 'train.py'
  colpyfile = '''# train.py for unitest
  import os
  env = os.environ.copy()
  assert "PADDLE_MASTER" in env
  assert "PADDLE_GLOBAL_SIZE" in env
  assert "PADDLE_LOCAL_SIZE" in env
  assert "PADDLE_GLOBAL_RANK" in env
  assert "PADDLE_LOCAL_RANK" in env
  '''
  
  pspyfile = '''# train.py for unitest
  import os
  env = os.environ.copy()
  assert "PADDLE_PSERVERS_IP_PORT_LIST" in env
  assert "PADDLE_TRAINER_ENDPOINTS" in env
  assert "PADDLE_ROLE" in env
  #assert "PADDLE_RANK" in env
  '''
  
  
  def write_file(name, ct):
      with open(name, "w") as f:
          f.write(ct)
  
  
  def get_files(pth, prefix):
      return [
          f for f in listdir(pth) if isfile(join(pth, f)) and f.startswith(prefix)
          and f != f"{prefix}.gpu.log"
      ]
  
  
  class Collective_Test(unittest.TestCase):
  
      def setUp(self):
          write_file(pyname, colpyfile)
  
      def pdrun(self, args, env=None):
          cmd = [sys.executable.split('/')[-1], "-m", "paddle.distributed.launch"]
          if args:
              cmd.extend(args.split(" "))
          cmd.extend([pyname])
          env = os.environ.copy()
          # virtual devies for testing
          env.update({'CUDA_VISIBLE_DEVICES': '0,1,2,3,4,5,6,7'})
          proc = subprocess.Popen(cmd, env=env)
          return proc
  
      def test_collective_1(self):
          args = "--job_id test1"
          p = self.pdrun(args)
          p.wait()
          self.assertTrue(p.poll() == 0)
  
      def test_collective_2(self):
          if os.path.exists('./log'):
              shutil.rmtree('./log')
  
          args = "--job_id test2 --devices 0,1,2"
          p = self.pdrun(args)
          p.wait()
          self.assertTrue(p.poll() == 0)
  
          c = get_files('log', 'test2')
          self.assertTrue(len(c) == 4)
  
      def test_collective_3(self):
          if os.path.exists('./log'):
              shutil.rmtree('./log')
  
          port = random.randrange(6000, 8000)
          args = "--job_id test3 --devices 0,1 --master 127.0.0.1:{} --nnodes 2".format(
              port)
          p1 = self.pdrun(args)
          p2 = self.pdrun(args)
          p1.wait()
          p2.wait()
          self.assertTrue(p1.poll() == 0)
          self.assertTrue(p2.poll() == 0)
  
          c = get_files('log', 'test3')
          self.assertTrue(len(c) == 6)
  
  
  class PS_Test(unittest.TestCase):
  
      def setUp(self):
          write_file(pyname, pspyfile)
  
      def pdrun(self, args, env=None):
          cmd = [sys.executable.split('/')[-1], "-m", "paddle.distributed.launch"]
          if args:
              cmd.extend(args.split(" "))
          cmd.extend([pyname])
          proc = subprocess.Popen(cmd, env)
          return proc
  
      def test_ps_1(self):
          args = "--run_mode ps"
          p = self.pdrun(args)
          p.wait()
          self.assertTrue(p.poll() == 0)
  
      def test_ps_2(self):
          if os.path.exists('./log'):
              shutil.rmtree('./log')
  
          args = "--job_id ps2 --server_num=2 --trainer_num=2"
          p = self.pdrun(args)
          p.wait()
          self.assertTrue(p.poll() == 0)
  
          c = get_files('log', 'ps2')
          self.assertTrue(len(c) == 5)
  
      def test_ps_3(self):
          if os.path.exists('./log'):
              shutil.rmtree('./log')
  
          port = random.randrange(6000, 8000)
          args = "--job_id ps3 --master 127.0.0.1:{} --nnodes 2 --server_num=1 --trainer_num=1".format(
              port)
          p1 = self.pdrun(args)
          p2 = self.pdrun(args)
          p1.wait()
          p2.wait()
          self.assertTrue(p1.poll() == 0)
          self.assertTrue(p2.poll() == 0)
  
          c = get_files('log', 'ps3')
          self.assertTrue(len(c) == 6)
  
      def test_ps_4(self):
          if os.path.exists('./log'):
              shutil.rmtree('./log')
  
          args = "--job_id ps4 --servers 127.0.0.1:8900,127.0.0.1:8901 --trainers 127.0.0.1:8902,127.0.0.1:8903"
          p1 = self.pdrun(args)
          p1.wait()
          self.assertTrue(p1.poll() == 0)
  
          c = get_files('log', 'ps4')
          self.assertTrue(len(c) == 5)
  
  
  if __name__ == '__main__':
      unittest.main()
  