FlagScale/tests/unit_tests/runner/test_launcher_ssh_perf.py at 95c7bc354c742b17fc623687452c96b767a9a41b · flagos-ai/FlagScale · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
from omegaconf import OmegaConf

from flagscale.runner.launcher.launcher_ssh import _get_runner_cmd_train


def test_get_runner_cmd_train_strips_perf_monitor_runner_keys():
    config = OmegaConf.create(
        {
            "experiment": {
                "runner": {
                    "backend": "torchrun",
                    "nnodes": 1,
                    "nproc_per_node": 8,
                    "rdzv_backend": "static",
                    "enable_perf_monitor": True,
                    "perf_log_interval": 5,
                    "perf_log_dir": "/tmp/perf_monitor",
                    "perf_console_output": True,
                }
            },
            "train": {
                "system": {
                    "logging": {
                        "details_dir": "/tmp/details",
                    }
                }
            },
        }
    )

    cmd = _get_runner_cmd_train("localhost", "127.0.0.1", 29500, 1, 0, 8, config)

    assert cmd[0] == "torchrun"
    assert "--enable_perf_monitor" not in cmd
    assert "--perf_log_interval" not in cmd
    assert "--perf_log_dir" not in cmd
    assert "--perf_console_output" not in cmd
    assert "--log_dir" in cmd
    assert "--rdzv_endpoint" in cmd


def test_get_runner_cmd_train_strips_hipprof_runner_keys():
    config = OmegaConf.create(
        {
            "experiment": {
                "runner": {
                    "backend": "torchrun",
                    "nnodes": 1,
                    "nproc_per_node": 8,
                    "rdzv_backend": "static",
                    "hipprof_bin_path": "/opt/dtk/bin/hipprof",
                    "hipprof_output_dir": "/tmp/hipprof",
                }
            },
            "train": {
                "system": {
                    "logging": {
                        "details_dir": "/tmp/details",
                    }
                },
                "model": {
                    "profile": True,
                    "use_hipprof_profiler": True,
                },
            },
        }
    )

    cmd = _get_runner_cmd_train("localhost", "127.0.0.1", 29500, 1, 0, 8, config)

    assert "--hipprof_bin_path" not in cmd
    assert "--hipprof_output_dir" not in cmd