diff --git a/test_runner.py b/test_runner.py index f7a4c7a4..dfd4a987 100755 --- a/test_runner.py +++ b/test_runner.py @@ -11,6 +11,8 @@ from dataclasses import dataclass from typing import Sequence +from torchtitan.logging_utils import logger + try: import tomllib except ModuleNotFoundError: @@ -191,7 +193,7 @@ def run_test(test_flavor: OverrideDefinitions, full_path: str): cmd = f"CONFIG_FILE={full_path} NGPU={test_flavor.ngpu} LOG_RANK=0,1,2,3 ./run_llama_train.sh" if override_arg: cmd += " " + " ".join(override_arg) - print( + logger.info( f"=====Integration test, flavor : {test_flavor.test_descr}, command : {cmd}=====" ) @@ -203,14 +205,14 @@ def run_test(test_flavor: OverrideDefinitions, full_path: str): assert ( dump_folder_arg is not None ), "Can't use seed checkpoint if folder is not specified" - print("Creating seed checkpoint") + logger.info("Creating seed checkpoint") result = _run_cmd( f"CONFIG_FILE={full_path} ./create_seed_checkpoint.sh {dump_folder_arg}" ) - print(result.stdout) + logger.info(result.stdout) result = _run_cmd(cmd) - print(result.stdout) + logger.info(result.stdout) if result.returncode != 0: raise Exception( f"Integration test failed, flavor : {test_flavor.test_descr}, command : {cmd}"