wehub-resource-sync
118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
# 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 contextlib
|
|
import itertools
|
|
import os
|
|
import shutil
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
import paddle
|
|
|
|
|
|
class CommunicationTestDistBase(unittest.TestCase):
|
|
def setUp(self, save_log_dir=None, num_of_devices=2, timeout=120, nnode=1):
|
|
if num_of_devices > paddle.device.cuda.device_count():
|
|
self.skipTest("number of GPUs is not enough")
|
|
|
|
self._python_interp = sys.executable
|
|
self._save_log_dir = save_log_dir
|
|
self._log_dir = tempfile.TemporaryDirectory()
|
|
self._num_of_devices = num_of_devices
|
|
self._device_list = [str(i) for i in range(num_of_devices)]
|
|
self._timeout = timeout
|
|
self._seeds = [i + 10 for i in range(num_of_devices)]
|
|
self._devices = ','.join(self._device_list)
|
|
self._nnode = nnode
|
|
self._port_set = set()
|
|
|
|
def _find_free_port(self):
|
|
def __free_port():
|
|
with contextlib.closing(
|
|
socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
) as s:
|
|
s.bind(('', 0))
|
|
return s.getsockname()[1]
|
|
|
|
while True:
|
|
port = __free_port()
|
|
if port not in self._port_set:
|
|
self._port_set.add(port)
|
|
return port
|
|
|
|
def run_test_case(self, script_file, user_defined_envs=None):
|
|
runtime_envs = os.environ
|
|
if user_defined_envs is not None:
|
|
runtime_envs.update(user_defined_envs)
|
|
runtime_envs["CUDA_VISIBLE_DEVICES"] = self._devices
|
|
if self._nnode > 1:
|
|
start_command = f"{self._python_interp} -u -m paddle.distributed.launch --nnode={self._nnode} --master=127.0.0.1:{self._find_free_port()} --log_dir {self._log_dir.name} --devices {self._devices} {script_file}"
|
|
else:
|
|
start_command = f"{self._python_interp} -u -m paddle.distributed.launch --log_dir {self._log_dir.name} --devices {self._devices} {script_file}"
|
|
start_command_list = start_command.strip().split()
|
|
|
|
if self._nnode > 1:
|
|
for i in range(1, self._nnode):
|
|
p = subprocess.Popen(
|
|
start_command_list,
|
|
env=runtime_envs,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
|
|
try:
|
|
self._launcher = subprocess.run(
|
|
start_command_list,
|
|
env=runtime_envs,
|
|
timeout=self._timeout,
|
|
check=True,
|
|
)
|
|
except subprocess.TimeoutExpired as err:
|
|
raise TimeoutError(
|
|
f"Timeout while running command {err.cmd}, try to set a longer period, {err.timeout} is not enough."
|
|
)
|
|
except subprocess.CalledProcessError as err:
|
|
raise RuntimeError(
|
|
f"Error occurs when running this test case. The return code of command {err.cmd} is {err.returncode}"
|
|
)
|
|
|
|
def tearDown(self):
|
|
if self._save_log_dir:
|
|
temp_log_dir_name = os.path.basename(self._log_dir.name)
|
|
dir_name = os.path.join(self._save_log_dir, temp_log_dir_name)
|
|
if not os.path.isdir(dir_name):
|
|
print(f"The running logs will copy to {dir_name}")
|
|
shutil.copytree(self._log_dir.name, dir_name)
|
|
else:
|
|
raise RuntimeError(
|
|
f"Directory {dir_name} exists, failed to save log."
|
|
)
|
|
|
|
|
|
def gen_product_envs_list(default_envs, changeable_envs):
|
|
envs_list = []
|
|
for values in itertools.product(*changeable_envs.values()):
|
|
envs = dict(zip(changeable_envs.keys(), values))
|
|
envs.update(default_envs)
|
|
envs_list.append(envs)
|
|
return envs_list
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|