run.py 18.8 KB
Newer Older
T
tangwei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# Copyright (c) 2020 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.

T
tangwei 已提交
15
import os
T
tangwei 已提交
16
import subprocess
X
test  
xjqbest 已提交
17
import sys
T
tangwei 已提交
18 19
import argparse
import tempfile
20
import warnings
C
Chengmo 已提交
21

X
fix  
xjqbest 已提交
22
import copy
23 24 25
from paddlerec.core.factory import TrainerFactory
from paddlerec.core.utils import envs
from paddlerec.core.utils import util
X
test  
xjqbest 已提交
26
from paddlerec.core.utils import validation
T
tangwei 已提交
27

T
tangwei 已提交
28 29
engines = {}
device = ["CPU", "GPU"]
J
Jinhua Liang 已提交
30 31

engine_choices = ["TRAIN", "INFER", "LOCAL_CLUSTER_TRAIN", "CLUSTER_TRAIN"]
T
tangwei 已提交
32 33


T
tangwei 已提交
34
def engine_registry():
T
tangwei 已提交
35 36 37
    engines["TRANSPILER"] = {}
    engines["PSLIB"] = {}

C
Chengmo 已提交
38 39 40
    engines["TRANSPILER"]["TRAIN"] = single_train_engine
    engines["TRANSPILER"]["INFER"] = single_infer_engine
    engines["TRANSPILER"]["LOCAL_CLUSTER_TRAIN"] = local_cluster_engine
T
tangwei 已提交
41
    engines["TRANSPILER"]["CLUSTER"] = cluster_engine
C
Chengmo 已提交
42 43 44
    engines["PSLIB"]["TRAIN"] = local_mpi_engine
    engines["PSLIB"]["LOCAL_CLUSTER_TRAIN"] = local_mpi_engine
    engines["PSLIB"]["CLUSTER_TRAIN"] = cluster_mpi_engine
T
tangwei 已提交
45
    engines["PSLIB"]["CLUSTER"] = cluster_mpi_engine
T
tangwei 已提交
46

T
tangwei 已提交
47

X
fix  
xjqbest 已提交
48
def get_inters_from_yaml(file, filters):
X
test  
xjqbest 已提交
49
    _envs = envs.load_yaml(file)
T
tangwei 已提交
50 51 52
    flattens = envs.flatten_environs(_envs)
    inters = {}
    for k, v in flattens.items():
X
fix  
xjqbest 已提交
53 54 55
        for f in filters:
            if k.startswith(f):
                inters[k] = v
T
tangwei 已提交
56
    return inters
T
tangwei 已提交
57 58


X
fix  
xjqbest 已提交
59
def get_all_inters_from_yaml(file, filters):
C
Chengmo 已提交
60
    _envs = envs.load_yaml(file)
X
fix  
xjqbest 已提交
61 62 63 64 65 66 67 68 69 70 71 72
    all_flattens = {}

    def fatten_env_namespace(namespace_nests, local_envs):
        for k, v in local_envs.items():
            if isinstance(v, dict):
                nests = copy.deepcopy(namespace_nests)
                nests.append(k)
                fatten_env_namespace(nests, v)
            elif (k == "dataset" or k == "phase" or
                  k == "runner") and isinstance(v, list):
                for i in v:
                    if i.get("name") is None:
C
Chengmo 已提交
73
                        raise ValueError("name must be in dataset list. ", v)
X
fix  
xjqbest 已提交
74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90
                    nests = copy.deepcopy(namespace_nests)
                    nests.append(k)
                    nests.append(i["name"])
                    fatten_env_namespace(nests, i)
            else:
                global_k = ".".join(namespace_nests + [k])
                all_flattens[global_k] = v

    fatten_env_namespace([], _envs)
    ret = {}
    for k, v in all_flattens.items():
        for f in filters:
            if k.startswith(f):
                ret[k] = v
    return ret


T
tangwei 已提交
91 92 93 94 95 96 97 98
def get_modes(running_config):
    if not isinstance(running_config, dict):
        raise ValueError("get_modes arguments must be [dict]")

    modes = running_config.get("mode")
    if not modes:
        raise ValueError("yaml mast have config: mode")

T
tangwei 已提交
99 100 101
    if isinstance(modes, str):
        modes = [modes]

T
tangwei 已提交
102 103 104 105
    return modes


def get_engine(args, running_config, mode):
T
tangwei 已提交
106
    transpiler = get_transpiler()
X
fix  
xjqbest 已提交
107

T
tangwei 已提交
108 109
    engine_class = ".".join(["runner", mode, "class"])
    engine_device = ".".join(["runner", mode, "device"])
C
Chengmo 已提交
110
    device_gpu_choices = ".".join(["runner", mode, "selected_gpus"])
T
tangwei 已提交
111 112

    engine = running_config.get(engine_class, None)
X
fix  
xjqbest 已提交
113
    if engine is None:
T
tangwei 已提交
114 115 116 117
        raise ValueError("not find {} in yaml, please check".format(
            mode, engine_class))
    device = running_config.get(engine_device, None)

T
tangwei 已提交
118 119 120
    engine = engine.upper()
    device = device.upper()

T
tangwei 已提交
121 122 123
    if device is None:
        print("not find device be specified in yaml, set CPU as default")
        device = "CPU"
C
Chengmo 已提交
124

T
tangwei 已提交
125
    if device == "GPU":
T
tangwei 已提交
126 127 128 129 130 131
        selected_gpus = running_config.get(device_gpu_choices, None)

        if selected_gpus is None:
            print(
                "not find selected_gpus be specified in yaml, set `0` as default"
            )
C
Chengmo 已提交
132
            selected_gpus = "0"
T
tangwei 已提交
133 134 135 136
        else:
            print("selected_gpus {} will be specified for running".format(
                selected_gpus))

C
Chengmo 已提交
137 138
        selected_gpus_num = len(selected_gpus.split(","))
        if selected_gpus_num > 1:
J
Jinhua Liang 已提交
139
            engine = "LOCAL_CLUSTER_TRAIN"
C
Chengmo 已提交
140

T
tangwei 已提交
141
    if engine not in engine_choices:
T
tangwei 已提交
142 143
        raise ValueError("{} can not be chosen in {}".format(engine_class,
                                                             engine_choices))
T
tangwei 已提交
144

T
tangwei 已提交
145
    run_engine = engines[transpiler].get(engine, None)
T
tangwei 已提交
146 147 148 149
    return run_engine


def get_transpiler():
T
tangwei 已提交
150
    FNULL = open(os.devnull, 'w')
T
tangwei 已提交
151 152 153 154
    cmd = [
        "python", "-c",
        "import paddle.fluid as fluid; fleet_ptr = fluid.core.Fleet(); [fleet_ptr.copy_table_by_feasign(10, 10, [2020, 1010])];"
    ]
T
tangwei 已提交
155 156 157
    proc = subprocess.Popen(cmd, stdout=FNULL, stderr=FNULL, cwd=os.getcwd())
    ret = proc.wait()
    if ret == -11:
T
tangwei 已提交
158
        return "PSLIB"
T
tangwei 已提交
159
    else:
T
tangwei 已提交
160
        return "TRANSPILER"
T
tangwei 已提交
161 162


T
tangwei 已提交
163 164 165
def set_runtime_envs(cluster_envs, engine_yaml):
    if cluster_envs is None:
        cluster_envs = {}
T
tangwei 已提交
166 167

    envs.set_runtime_environs(cluster_envs)
T
fix bug  
tangwei 已提交
168 169 170

    need_print = {}
    for k, v in os.environ.items():
T
tangwei 已提交
171
        if k.startswith("train.trainer."):
T
fix bug  
tangwei 已提交
172 173 174
            need_print[k] = v

    print(envs.pretty_print_envs(need_print, ("Runtime Envs", "Value")))
T
tangwei 已提交
175 176


C
Chengmo 已提交
177
def single_train_engine(args):
T
tangwei 已提交
178 179 180 181 182 183 184 185 186 187 188 189
    run_extras = get_all_inters_from_yaml(args.model, ["runner."])
    mode = envs.get_runtime_environ("mode")
    trainer_class = ".".join(["runner", mode, "trainer_class"])
    fleet_class = ".".join(["runner", mode, "fleet_mode"])
    device_class = ".".join(["runner", mode, "device"])
    selected_gpus_class = ".".join(["runner", mode, "selected_gpus"])

    trainer = run_extras.get(trainer_class, "GeneralTrainer")
    fleet_mode = run_extras.get(fleet_class, "ps")
    device = run_extras.get(device_class, "cpu")
    selected_gpus = run_extras.get(selected_gpus_class, "0")
    executor_mode = "train"
T
tangwei 已提交
190

T
tangwei 已提交
191
    single_envs = {}
C
Chengmo 已提交
192 193

    if device.upper() == "GPU":
T
tangwei 已提交
194 195 196 197 198 199 200 201
        selected_gpus_num = len(selected_gpus.split(","))
        if selected_gpus_num != 1:
            raise ValueError(
                "Single Mode Only Support One GPU, Set Local Cluster Mode to use Multi-GPUS"
            )

        single_envs["selsected_gpus"] = selected_gpus
        single_envs["FLAGS_selected_gpus"] = selected_gpus
C
chengmo 已提交
202

C
chengmo 已提交
203
    single_envs["train.trainer.trainer"] = trainer
C
Chengmo 已提交
204 205
    single_envs["fleet_mode"] = fleet_mode
    single_envs["train.trainer.executor_mode"] = executor_mode
C
chengmo 已提交
206 207
    single_envs["train.trainer.threads"] = "2"
    single_envs["train.trainer.platform"] = envs.get_platform()
C
Chengmo 已提交
208 209
    single_envs["train.trainer.engine"] = "single"

X
fix  
xjqbest 已提交
210 211 212
    set_runtime_envs(single_envs, args.model)
    trainer = TrainerFactory.create(args.model)
    return trainer
X
fix  
xjqbest 已提交
213

X
fix  
xjqbest 已提交
214 215

def single_infer_engine(args):
T
tangwei 已提交
216
    run_extras = get_all_inters_from_yaml(args.model, ["runner."])
C
Chengmo 已提交
217

T
tangwei 已提交
218 219 220 221 222
    mode = envs.get_runtime_environ("mode")
    trainer_class = ".".join(["runner", mode, "trainer_class"])
    fleet_class = ".".join(["runner", mode, "fleet_mode"])
    device_class = ".".join(["runner", mode, "device"])
    selected_gpus_class = ".".join(["runner", mode, "selected_gpus"])
C
Chengmo 已提交
223

224 225 226 227 228 229 230 231
    epochs_class = ".".join(["runner", mode, "epochs"])
    epochs = run_extras.get(epochs_class, 1)
    if epochs > 1:
        warnings.warn(
            "It makes no sense to predict the same model for multiple epochs",
            category=UserWarning,
            stacklevel=2)

T
tangwei 已提交
232 233 234 235
    trainer = run_extras.get(trainer_class, "GeneralTrainer")
    fleet_mode = run_extras.get(fleet_class, "ps")
    device = run_extras.get(device_class, "cpu")
    selected_gpus = run_extras.get(selected_gpus_class, "0")
C
Chengmo 已提交
236 237
    executor_mode = "infer"

T
tangwei 已提交
238 239
    single_envs = {}

C
Chengmo 已提交
240
    if device.upper() == "GPU":
T
tangwei 已提交
241 242 243 244 245 246 247 248
        selected_gpus_num = len(selected_gpus.split(","))
        if selected_gpus_num != 1:
            raise ValueError(
                "Single Mode Only Support One GPU, Set Local Cluster Mode to use Multi-GPUS"
            )

        single_envs["selsected_gpus"] = selected_gpus
        single_envs["FLAGS_selected_gpus"] = selected_gpus
C
Chengmo 已提交
249

X
fix  
xjqbest 已提交
250
    single_envs["train.trainer.trainer"] = trainer
C
Chengmo 已提交
251 252
    single_envs["train.trainer.executor_mode"] = executor_mode
    single_envs["fleet_mode"] = fleet_mode
X
fix  
xjqbest 已提交
253 254
    single_envs["train.trainer.threads"] = "2"
    single_envs["train.trainer.platform"] = envs.get_platform()
C
Chengmo 已提交
255 256
    single_envs["train.trainer.engine"] = "single"

X
fix  
xjqbest 已提交
257 258 259
    set_runtime_envs(single_envs, args.model)
    trainer = TrainerFactory.create(args.model)
    return trainer
C
chengmo 已提交
260

X
fix  
xjqbest 已提交
261

T
tangwei 已提交
262
def cluster_engine(args):
T
tangwei 已提交
263
    def master():
264
        from paddlerec.core.engine.cluster.cluster import ClusterEngine
X
test  
xjqbest 已提交
265
        _envs = envs.load_yaml(args.backend)
T
tangwei 已提交
266
        flattens = envs.flatten_environs(_envs, "_")
J
Jinhua Liang 已提交
267 268
        flattens["engine_role"] = "MASTER"
        flattens["engine_mode"] = envs.get_runtime_environ("mode")
T
tangwei 已提交
269
        flattens["engine_run_config"] = args.model
T
tangwei 已提交
270 271
        flattens["engine_temp_path"] = tempfile.mkdtemp()
        envs.set_runtime_environs(flattens)
J
Jinhua Liang 已提交
272
        ClusterEngine.workspace_replace()
C
Chengmo 已提交
273
        print(envs.pretty_print_envs(flattens, ("Submit Envs", "Value")))
T
tangwei 已提交
274 275 276 277

        launch = ClusterEngine(None, args.model)
        return launch

J
Jinhua Liang 已提交
278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299
    def worker(mode):
        if not mode:
            raise ValueError("mode: {} can not be recognized")

        run_extras = get_all_inters_from_yaml(args.model, ["runner."])

        trainer_class = ".".join(["runner", mode, "trainer_class"])
        fleet_class = ".".join(["runner", mode, "fleet_mode"])
        device_class = ".".join(["runner", mode, "device"])
        selected_gpus_class = ".".join(["runner", mode, "selected_gpus"])
        strategy_class = ".".join(["runner", mode, "distribute_strategy"])
        worker_class = ".".join(["runner", mode, "worker_num"])
        server_class = ".".join(["runner", mode, "server_num"])

        trainer = run_extras.get(trainer_class, "GeneralTrainer")
        fleet_mode = run_extras.get(fleet_class, "ps")
        device = run_extras.get(device_class, "cpu")
        selected_gpus = run_extras.get(selected_gpus_class, "0")
        distributed_strategy = run_extras.get(strategy_class, "async")
        worker_num = run_extras.get(worker_class, 1)
        server_num = run_extras.get(server_class, 1)
        executor_mode = "train"
C
Chengmo 已提交
300

J
Jinhua Liang 已提交
301 302
        device = device.upper()
        fleet_mode = fleet_mode.upper()
C
Chengmo 已提交
303

J
Jinhua Liang 已提交
304 305
        if fleet_mode == "COLLECTIVE" and device != "GPU":
            raise ValueError("COLLECTIVE can not be used with GPU")
C
Chengmo 已提交
306

J
Jinhua Liang 已提交
307
        cluster_envs = {}
C
Chengmo 已提交
308

J
Jinhua Liang 已提交
309 310
        if device == "GPU":
            cluster_envs["selected_gpus"] = selected_gpus
C
Chengmo 已提交
311

J
Jinhua Liang 已提交
312 313
        cluster_envs["server_num"] = server_num
        cluster_envs["worker_num"] = worker_num
C
Chengmo 已提交
314
        cluster_envs["fleet_mode"] = fleet_mode
T
tangwei 已提交
315 316
        cluster_envs["train.trainer.trainer"] = trainer
        cluster_envs["train.trainer.engine"] = "cluster"
J
Jinhua Liang 已提交
317
        cluster_envs["train.trainer.executor_mode"] = executor_mode
C
Chengmo 已提交
318
        cluster_envs["train.trainer.strategy"] = distributed_strategy
T
tangwei 已提交
319 320
        cluster_envs["train.trainer.threads"] = envs.get_runtime_environ(
            "CPU_NUM")
T
tangwei 已提交
321
        cluster_envs["train.trainer.platform"] = envs.get_platform()
C
chengmo 已提交
322 323
        print("launch {} engine with cluster to with model: {}".format(
            trainer, args.model))
T
tangwei 已提交
324
        set_runtime_envs(cluster_envs, args.model)
T
tangwei 已提交
325

T
bug fix  
tangwei12 已提交
326 327
        trainer = TrainerFactory.create(args.model)
        return trainer
T
tangwei 已提交
328

T
tangwei 已提交
329 330 331
    role = os.getenv("PADDLE_PADDLEREC_ROLE", "MASTER")

    if role == "WORKER":
J
Jinhua Liang 已提交
332 333
        mode = os.getenv("PADDLE_PADDLEREC_MODE", None)
        return worker(mode)
T
tangwei 已提交
334 335
    else:
        return master()
C
chengmo 已提交
336 337


T
tangwei 已提交
338
def cluster_mpi_engine(args):
T
tangwei 已提交
339 340
    print("launch cluster engine with cluster to run model: {}".format(
        args.model))
T
tangwei 已提交
341

T
fix bug  
tangwei 已提交
342
    cluster_envs = {}
T
tangwei 已提交
343
    cluster_envs["train.trainer.trainer"] = "CtrCodingTrainer"
T
tangwei 已提交
344
    cluster_envs["train.trainer.platform"] = envs.get_platform()
T
tangwei 已提交
345

T
tangwei 已提交
346
    set_runtime_envs(cluster_envs, args.model)
T
tangwei 已提交
347

T
tangwei 已提交
348 349 350 351 352
    trainer = TrainerFactory.create(args.model)
    return trainer


def local_cluster_engine(args):
J
Jinhua Liang 已提交
353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383
    def get_worker_num(run_extras, workers):
        _envs = envs.load_yaml(args.model)
        mode = envs.get_runtime_environ("mode")
        workspace = envs.get_runtime_environ("workspace")
        phases_class = ".".join(["runner", mode, "phases"])
        phase_names = run_extras.get(phases_class)
        phases = []
        all_phases = _envs.get("phase")
        if phase_names is None:
            phases = all_phases
        else:
            for phase in all_phases:
                if phase["name"] in phase_names:
                    phases.append(phase)

        dataset_names = []
        for phase in phases:
            dataset_names.append(phase["dataset_name"])

        datapaths = []
        for dataset in _envs.get("dataset"):
            if dataset["name"] in dataset_names:
                datapaths.append(dataset["data_path"])

        if not datapaths:
            raise ValueError("data path must exist for training/inference")

        datapaths = [
            envs.workspace_adapter_by_specific(path, workspace)
            for path in datapaths
        ]
384

J
Jinhua Liang 已提交
385 386
        all_workers = [len(os.listdir(path)) for path in datapaths]
        all_workers.append(workers)
387 388 389 390 391 392 393 394 395 396
        max_worker_num = min(all_workers)

        if max_worker_num >= workers:
            return workers

        print(
            "phases do not have enough datas for training, set worker/gpu cards num from {} to {}".
            format(workers, max_worker_num))

        return max_worker_num
C
chengmo 已提交
397

J
Jinhua Liang 已提交
398
    from paddlerec.core.engine.local_cluster import LocalClusterEngine
C
Chengmo 已提交
399

J
Jinhua Liang 已提交
400 401 402 403 404 405 406 407 408
    run_extras = get_all_inters_from_yaml(args.model, ["runner."])
    mode = envs.get_runtime_environ("mode")
    trainer_class = ".".join(["runner", mode, "trainer_class"])
    fleet_class = ".".join(["runner", mode, "fleet_mode"])
    device_class = ".".join(["runner", mode, "device"])
    selected_gpus_class = ".".join(["runner", mode, "selected_gpus"])
    strategy_class = ".".join(["runner", mode, "distribute_strategy"])
    worker_class = ".".join(["runner", mode, "worker_num"])
    server_class = ".".join(["runner", mode, "server_num"])
C
Chengmo 已提交
409

J
Jinhua Liang 已提交
410 411 412 413 414
    trainer = run_extras.get(trainer_class, "GeneralTrainer")
    fleet_mode = run_extras.get(fleet_class, "ps")
    device = run_extras.get(device_class, "cpu")
    selected_gpus = run_extras.get(selected_gpus_class, "0")
    distributed_strategy = run_extras.get(strategy_class, "async")
C
Chengmo 已提交
415
    executor_mode = "train"
J
Jinhua Liang 已提交
416 417 418 419 420 421 422

    worker_num = run_extras.get(worker_class, 1)
    server_num = run_extras.get(server_class, 1)

    device = device.upper()
    fleet_mode = fleet_mode.upper()

423 424 425 426 427 428 429 430 431 432
    cluster_envs = {}

    # Todo: delete follow hard code when paddle support ps-gpu.
    if device == "CPU":
        fleet_mode = "PS"
    elif device == "GPU":
        fleet_mode = "COLLECTIVE"
    if fleet_mode == "PS" and device != "CPU":
        raise ValueError("PS can not be used with GPU")

J
Jinhua Liang 已提交
433
    if fleet_mode == "COLLECTIVE" and device != "GPU":
434
        raise ValueError("COLLECTIVE can not be used without GPU")
C
Chengmo 已提交
435

436 437
    if fleet_mode == "PS":
        worker_num = get_worker_num(run_extras, worker_num)
J
Jinhua Liang 已提交
438

439
    if fleet_mode == "COLLECTIVE":
J
Jinhua Liang 已提交
440
        cluster_envs["selected_gpus"] = selected_gpus
441 442 443
        gpus = selected_gpus.split(",")
        gpu_num = get_worker_num(run_extras, len(gpus))
        cluster_envs["selected_gpus"] = ','.join(gpus[:gpu_num])
J
Jinhua Liang 已提交
444

C
Chengmo 已提交
445 446
    cluster_envs["server_num"] = server_num
    cluster_envs["worker_num"] = worker_num
C
chengmo 已提交
447
    cluster_envs["start_port"] = envs.find_free_port()
C
Chengmo 已提交
448
    cluster_envs["fleet_mode"] = fleet_mode
C
chengmo 已提交
449
    cluster_envs["log_dir"] = "logs"
C
chengmo 已提交
450
    cluster_envs["train.trainer.trainer"] = trainer
C
Chengmo 已提交
451 452
    cluster_envs["train.trainer.executor_mode"] = executor_mode
    cluster_envs["train.trainer.strategy"] = distributed_strategy
C
chengmo 已提交
453
    cluster_envs["train.trainer.threads"] = "2"
J
Jinhua Liang 已提交
454
    cluster_envs["CPU_NUM"] = cluster_envs["train.trainer.threads"]
C
chengmo 已提交
455 456 457
    cluster_envs["train.trainer.engine"] = "local_cluster"
    cluster_envs["train.trainer.platform"] = envs.get_platform()

T
tangwei 已提交
458 459
    print("launch {} engine with cluster to run model: {}".format(trainer,
                                                                  args.model))
C
chengmo 已提交
460 461 462 463 464 465

    set_runtime_envs(cluster_envs, args.model)
    launch = LocalClusterEngine(cluster_envs, args.model)
    return launch


T
tangwei 已提交
466
def local_mpi_engine(args):
T
tangwei 已提交
467 468
    print("launch cluster engine with cluster to run model: {}".format(
        args.model))
469
    from paddlerec.core.engine.local_mpi import LocalMPIEngine
T
tangwei 已提交
470

T
tangwei 已提交
471 472
    print("use 1X1 MPI ClusterTraining at localhost to run model: {}".format(
        args.model))
T
tangwei 已提交
473

T
tangwei 已提交
474 475 476
    mpi = util.run_which("mpirun")
    if not mpi:
        raise RuntimeError("can not find mpirun, please check environment")
C
Chengmo 已提交
477

J
Jinhua Liang 已提交
478 479 480 481 482 483
    run_extras = get_all_inters_from_yaml(args.model, ["runner."])

    mode = envs.get_runtime_environ("mode")
    trainer_class = ".".join(["runner", mode, "trainer_class"])
    fleet_class = ".".join(["runner", mode, "fleet_mode"])
    distributed_strategy = "async"
C
Chengmo 已提交
484 485
    executor_mode = "train"

J
Jinhua Liang 已提交
486 487
    trainer = run_extras.get(trainer_class, "GeneralTrainer")
    fleet_mode = run_extras.get(fleet_class, "ps")
C
Chengmo 已提交
488

T
fix bug  
tangwei 已提交
489 490
    cluster_envs = {}
    cluster_envs["mpirun"] = mpi
C
Chengmo 已提交
491
    cluster_envs["train.trainer.trainer"] = trainer
T
fix bug  
tangwei 已提交
492
    cluster_envs["log_dir"] = "logs"
T
tangwei 已提交
493
    cluster_envs["train.trainer.engine"] = "local_cluster"
C
Chengmo 已提交
494 495 496 497
    cluster_envs["train.trainer.executor_mode"] = executor_mode
    cluster_envs["fleet_mode"] = fleet_mode
    cluster_envs["train.trainer.strategy"] = distributed_strategy
    cluster_envs["train.trainer.threads"] = "2"
T
tangwei 已提交
498
    cluster_envs["train.trainer.platform"] = envs.get_platform()
T
tangwei 已提交
499

T
tangwei 已提交
500
    set_runtime_envs(cluster_envs, args.model)
T
tangwei 已提交
501 502 503 504
    launch = LocalMPIEngine(cluster_envs, args.model)
    return launch


T
tangwei 已提交
505
def get_abs_model(model):
506
    if model.startswith("paddlerec."):
T
tangwei 已提交
507
        dir = envs.paddlerec_adapter(model)
T
tangwei 已提交
508
        path = os.path.join(dir, "config.yaml")
T
tangwei 已提交
509 510 511 512 513 514 515
    else:
        if not os.path.isfile(model):
            raise IOError("model config: {} invalid".format(model))
        path = model
    return path


T
tangwei 已提交
516
if __name__ == "__main__":
517
    parser = argparse.ArgumentParser(description='paddle-rec run')
T
tangwei 已提交
518
    parser.add_argument("-m", "--model", type=str)
T
tangwei 已提交
519
    parser.add_argument("-b", "--backend", type=str, default=None)
T
tangwei 已提交
520

T
tangwei 已提交
521 522 523
    abs_dir = os.path.dirname(os.path.abspath(__file__))
    envs.set_runtime_environs({"PACKAGE_BASE": abs_dir})

T
tangwei 已提交
524
    args = parser.parse_args()
T
tangwei 已提交
525
    args.model = get_abs_model(args.model)
T
tangwei 已提交
526

X
test  
xjqbest 已提交
527 528
    if not validation.yaml_validation(args.model):
        sys.exit(-1)
T
tangwei 已提交
529

T
tangwei 已提交
530
    engine_registry()
J
Jinhua Liang 已提交
531 532
    running_config = get_all_inters_from_yaml(
        args.model, ["workspace", "mode", "runner."])
T
tangwei 已提交
533 534 535
    modes = get_modes(running_config)

    for mode in modes:
J
Jinhua Liang 已提交
536 537 538 539
        envs.set_runtime_environs({
            "mode": mode,
            "workspace": running_config["workspace"]
        })
T
tangwei 已提交
540 541 542
        which_engine = get_engine(args, running_config, mode)
        engine = which_engine(args)
        engine.run()