wehub-resource-sync
103 lines
3.4 KiB
Python
103 lines
3.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import subprocess
|
|
|
|
from types import SimpleNamespace
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.launcher.multinode_runner import MultiNodeRunner
|
|
|
|
|
|
class DummyRunner(MultiNodeRunner):
|
|
|
|
def backend_exists(self):
|
|
return True
|
|
|
|
def get_cmd(self, environment, active_resources):
|
|
return []
|
|
|
|
|
|
if not get_accelerator().is_available():
|
|
pytest.skip("only supported in accelerator environments.", allow_module_level=True)
|
|
|
|
user_arg_test_script = """import argparse
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--prompt", type=str)
|
|
parser.add_argument("--local_rank", type=int, default=0)
|
|
parser.add_argument("--world_size", type=int, default=1)
|
|
args = parser.parse_args()
|
|
print("ARG PARSE SUCCESS")
|
|
"""
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def user_script_fp(tmpdir):
|
|
script_fp = tmpdir.join("user_arg_test.py")
|
|
with open(script_fp, "w") as f:
|
|
f.write(user_arg_test_script)
|
|
return script_fp
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def cmd(user_script_fp, prompt, multi_node):
|
|
if multi_node:
|
|
cmd = ("deepspeed", "--force_multi", "--num_nodes", "1", "--num_gpus", "1", user_script_fp, "--prompt", prompt)
|
|
else:
|
|
cmd = ("deepspeed", "--num_nodes", "1", "--num_gpus", "1", user_script_fp, "--prompt", prompt)
|
|
return cmd
|
|
|
|
|
|
@pytest.fixture
|
|
def dummy_runner():
|
|
args = SimpleNamespace(user_args=[], user_script="dummy_script.py")
|
|
return DummyRunner(args, "dummy_world_info")
|
|
|
|
|
|
@pytest.mark.parametrize("prompt", [
|
|
'''"I am 6' tall"''', """'I am 72" tall'""", """'"translate English to Romanian: "'""",
|
|
'''I'm going to tell them "DeepSpeed is the best"'''
|
|
])
|
|
@pytest.mark.parametrize("multi_node", [True, False])
|
|
def test_user_args(cmd, multi_node):
|
|
if multi_node and get_accelerator().device_name() == "cpu":
|
|
pytest.skip("CPU accelerator does not support this test yet")
|
|
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
out, err = p.communicate()
|
|
assert "ARG PARSE SUCCESS" in out.decode("utf-8"), f"User args not parsed correctly: {err.decode('utf-8')}"
|
|
|
|
|
|
def test_bash_string_args(tmpdir, user_script_fp):
|
|
bash_script = f"""
|
|
ARGS="--prompt 'DeepSpeed is the best'"
|
|
echo ${{ARGS}}|xargs deepspeed --num_nodes 1 --num_gpus 1 {user_script_fp}
|
|
"""
|
|
|
|
bash_fp = tmpdir.join("bash_script.sh")
|
|
with open(bash_fp, "w") as f:
|
|
f.write(bash_script)
|
|
|
|
p = subprocess.Popen(["bash", bash_fp], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
out, err = p.communicate()
|
|
assert "ARG PARSE SUCCESS" in out.decode("utf-8"), f"User args not parsed correctly: {err.decode('utf-8')}"
|
|
|
|
|
|
def test_add_export_with_special_characters(dummy_runner):
|
|
"""
|
|
Values with special characters (e.g., 64(x2)) must be quoted to avoid bash syntax errors.
|
|
"""
|
|
dummy_runner.add_export("SLURM_JOB_CPUS_PER_NODE", "64(x2)")
|
|
assert dummy_runner.exports["SLURM_JOB_CPUS_PER_NODE"] == "\"64(x2)\""
|
|
|
|
|
|
def test_add_export_no_special_characters(dummy_runner):
|
|
"""
|
|
Values without special characters should remain unquoted (e.g., PYTHONPATH).
|
|
This avoids issues where unnecessary quotes break module imports.
|
|
"""
|
|
dummy_runner.add_export("PYTHONPATH", "/usr/local/lib/python3.9/site-packages")
|
|
assert dummy_runner.exports["PYTHONPATH"] == "/usr/local/lib/python3.9/site-packages"
|