code: def main():
      check_pip_packages()
  
      options = parse_args()
  
      # Include sharding info in all metrics
      which_shard, num_shards = get_sharding_opts(options)
      add_global_metric("shard", which_shard)
      add_global_metric("num_shards", num_shards)
  
      test_directory = str(REPO_ROOT / "test")
      selected_tests = get_selected_tests(options)
  
      os.makedirs(REPO_ROOT / "test" / "test-reports", exist_ok=True)
  
      if options.coverage and not PYTORCH_COLLECT_COVERAGE:
          shell(["coverage", "erase"])
  
      aggregated_heuristics: AggregatedHeuristics = (
          AggregatedHeuristics(unranked_tests=selected_tests)
      )
  
      with open(
          REPO_ROOT
          / "test"
          / "test-reports"
          / "td_heuristic_rankings.log",
          "w",
      ) as f:
          if IS_CI:
              # downloading test cases configuration to local environment
              get_test_case_configs(dirpath=test_directory)
              aggregated_heuristics = get_test_prioritizations(
                  selected_tests, file=f
              )
  
          test_prioritizations = (
              aggregated_heuristics.get_aggregated_priorities()
          )
  
          f.write(test_prioritizations.get_info_str())
  
      test_file_times_dict = load_test_file_times()
      test_class_times_dict = load_test_class_times()
  
      class TestBatch:
          """Defines a set of tests with similar priority that should be run together on the current shard"""
  
          name: str
          sharded_tests: List[ShardedTest]
          failures: List[TestFailure]
  
          def __init__(
              self,
              name: str,
              raw_tests: Sequence[TestRun],
              should_sort_shard: bool,
          ):
              self.name = name
              self.failures = []
              self.sharded_tests = do_sharding(
                  options,
                  raw_tests,
                  test_file_times_dict,
                  test_class_times_dict,
                  sort_by_time=should_sort_shard,
              )
  
          def __str__(self):
              s = f"Name: {self.name}\n"
              s += "  Parallel tests:\n"
              s += "".join(
                  f"    {test}\n"
                  for test in self.sharded_tests
                  if not must_serial(test)
              )
              s += "  Serial tests:\n"
              s += "".join(
                  f"    {test}\n"
                  for test in self.sharded_tests
                  if must_serial(test)
              )
              return s.strip()
  
      test_batches: List[TestBatch] = []
  
      # Each batch will be run sequentially
      test_batches = [
          TestBatch(
              "high_relevance",
              test_prioritizations.get_high_relevance_tests(),
              False,
          ),
          TestBatch(
              "probable_relevance",
              test_prioritizations.get_probable_relevance_tests(),
              False,
          ),
          TestBatch(
              "unranked_relevance",
              test_prioritizations.get_unranked_relevance_tests(),
              True,
          ),
          TestBatch(
              "unlikely_relevance",
              test_prioritizations.get_unlikely_relevance_tests(),
              True,
          ),
          TestBatch(
              "none_relevance",
              test_prioritizations.get_none_relevance_tests(),
              True,
          ),
      ]
  
      for test_batch in test_batches:
          print_to_stderr(test_batch)
  
      if options.dry_run:
          return
  
      if options.dynamo:
          os.environ["PYTORCH_TEST_WITH_DYNAMO"] = "1"
      elif options.inductor:
          os.environ["PYTORCH_TEST_WITH_INDUCTOR"] = "1"
  
      if not options.no_translation_validation:
          os.environ["PYTORCH_TEST_WITH_TV"] = "1"
  
      try:
          # Actually run the tests
          start_time = time.time()
          for test_batch in test_batches:
              elapsed_time = time.time() - start_time
              print_to_stderr(
                  f"Starting test batch '{test_batch.name}' {elapsed_time} seconds after initiating testing"
              )
              print_to_stderr(
                  f"With sharding, this batch will run {len(test_batch.sharded_tests)} tests"
              )
              run_tests(
                  test_batch.sharded_tests,
                  test_directory,
                  options,
                  test_batch.failures,
              )
  
      finally:
          if options.coverage:
              from coverage import Coverage
  
              with set_cwd(test_directory):
                  cov = Coverage()
                  if PYTORCH_COLLECT_COVERAGE:
                      cov.load()
                  cov.combine(strict=False)
                  cov.save()
                  if not PYTORCH_COLLECT_COVERAGE:
                      cov.html_report()
  
          all_failures = [
              failure
              for batch in test_batches
              for failure in batch.failures
          ]
  
          if IS_CI:
              num_tests = len(selected_tests)
              for test, _ in all_failures:
                  test_stats = aggregated_heuristics.get_test_stats(
                      test
                  )
                  test_stats["num_total_tests"] = num_tests
  
                  print_to_stderr("Emiting td_test_failure_stats")
                  emit_metric(
                      "td_test_failure_stats",
                      {
                          **test_stats,
                          "confidence_ratings": get_prediction_confidences(
                              selected_tests
                          ),
                          "failure": str(test),
                          "tests": selected_tests,
                      },
                  )
  
      if len(all_failures):
          for _, err in all_failures:
              print_to_stderr(err)
  
          # A disabled test is expected to fail, so there is no need to report a failure here
          if not RERUN_DISABLED_TESTS:
              sys.exit(1)
  
imports: import argparse
  import copy
  import glob
  import json
  import os
  import pathlib
  import re
  import shutil
  import signal
  import subprocess
  import sys
  import tempfile
  import time
  from collections import defaultdict
  from contextlib import ExitStack
  from datetime import datetime
  from typing import Any, cast, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union
  import pkg_resources
  import torch
  import torch.distributed
  from torch.multiprocessing import current_process, get_context
  from torch.testing._internal.common_utils import FILE_SCHEMA, get_report_path, IS_CI, parser, retry_shell, set_cwd, shell, TEST_WITH_ASAN, TEST_WITH_ROCM, TEST_WITH_SLOW_GRADCHECK
  from tools.stats.import_test_stats import ADDITIONAL_CI_FILES_FOLDER, TEST_CLASS_TIMES_FILE, TEST_TIMES_FILE
  from tools.stats.upload_metrics import add_global_metric, emit_metric
  from tools.testing.target_determination.determinator import AggregatedHeuristics, get_prediction_confidences, get_test_prioritizations
  from tools.testing.test_run import TestRun
  from tools.testing.test_selections import calculate_shards, get_test_case_configs, NUM_PROCS, ShardedTest, THRESHOLD
label_desc: Use image node_id 31 for calling a global function with example usage: main() and returns: s