main() — pytorch Function Reference

Architecture documentation for the main() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking calls 5 called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ main() — pytorch Function Reference

Dependency Diagram

graph TD
  60550900_f047_a85b_c5ab_ff80b30486b4["main()"]
  b93ae0ef_d85a_81fb_b07f_38119c7d07cc["huggingface_main()"]
  b93ae0ef_d85a_81fb_b07f_38119c7d07cc -->|calls| 60550900_f047_a85b_c5ab_ff80b30486b4
  0b66073b_10d9_d2e9_f83d_01d82f509823["parse_args()"]
  60550900_f047_a85b_c5ab_ff80b30486b4 -->|calls| 0b66073b_10d9_d2e9_f83d_01d82f509823
  43ce1a6e_7fe2_4da2_e9e2_87cf5c9d44a4["should_diff_branch()"]
  60550900_f047_a85b_c5ab_ff80b30486b4 -->|calls| 43ce1a6e_7fe2_4da2_e9e2_87cf5c9d44a4
  ce8fd365_4112_b289_9c73_7345d5e35203["RuntimeError()"]
  60550900_f047_a85b_c5ab_ff80b30486b4 -->|calls| ce8fd365_4112_b289_9c73_7345d5e35203
  06048c7c_b216_de58_82e5_19ecf074c96a["maybe_fresh_cache()"]
  60550900_f047_a85b_c5ab_ff80b30486b4 -->|calls| 06048c7c_b216_de58_82e5_19ecf074c96a
  27245bba_7bce_3446_86ed_29630ca8fe8a["process_entry()"]
  60550900_f047_a85b_c5ab_ff80b30486b4 -->|calls| 27245bba_7bce_3446_86ed_29630ca8fe8a
  style 60550900_f047_a85b_c5ab_ff80b30486b4 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 3752–3811

def main(runner, original_dir=None, args=None):
    if original_dir:
        os.chdir(original_dir)
    args = parse_args() if not args else parse_args(args)
    if args.baseline:
        args.baseline = os.path.abspath(args.baseline)

    if should_diff_branch(args):
        import git

        # We do this here so we error out earlier if there's an issue
        repo = git.Repo()
        if repo.is_dirty():
            raise RuntimeError(
                "--diff-branch called on dirty branch. Commit, stash, or reset."
            )
        main_branch = repo.active_branch.name
        if main_branch == args.diff_branch:
            raise RuntimeError(
                f"--diff-branch: current branch is same as {args.diff_branch} branch, what are you diffing?"
            )

    with maybe_fresh_cache(args):
        if args.caching_precompile:
            os.environ["TORCH_CACHING_PRECOMPILE"] = "1"
            torch._dynamo.config.caching_precompile = True

        args.init_distributed = args.only and args.multiprocess
        if args.init_distributed:
            # NB: Do NOT query device count before CUDA initialization; we're
            # going to overwrite CUDA_VISIBLE_DEVICES and this will result in
            # https://github.com/pytorch/pytorch/issues/107300
            device_count = torch.cuda.device_count()
            if device_count <= 1:
                log.warning(
                    "The use multiprocess flag is set but there are <= 1 devices available."
                )
            # multiprocess path
            args.world_size = device_count
            mp.spawn(
                process_entry, args=(runner, original_dir, args), nprocs=device_count
            )
        elif args.only and args.warm_start_latency:
            # Warm start mode. Enable FX graph caching and perform back-to-back runs in
            # separate processes (but ensure the inductor cache is preserved across runs).
            env = os.environ.copy()
            env["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "1"
            cmd = [sys.executable] + sys.argv
            cmd.remove("--warm-start-latency")

            print(f"Performing cold-start run for {args.only}")
            warmup_cmd = cmd + ["--repeat=1", "--disable-output"]
            subprocess.check_call(warmup_cmd, timeout=args.timeout, env=env)

            print(f"Performing warm-start run for {args.only}")
            subprocess.check_call(cmd, timeout=args.timeout, env=env)
        else:
            # single process path just uses the main process
            args.world_size = 1
            process_entry(0, runner, original_dir, args)