[test stats] use published test stats for sharding (#81116)

Use the nightly-published test stats to perform sharding, instead of
calculating it in every build job.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81116
Approved by: https://github.com/janeyx99

tools/stats/import_test_stats.py

@@ -41,6 +41,7 @@ def fetch_and_cache(
     This fetch and cache utils allows sharing between different process.
     """
     path = os.path.join(dirpath, name)
+    print(f"Downloading {url} to {path}")
 
     def is_cached_file_valid() -> bool:
         # Check if the file is new enough (see: FILE_CACHE_LIFESPAN_SECONDS). A real check
@@ -80,6 +81,21 @@ def get_slow_tests(
         return {}
 
 
+def get_test_times(dirpath: str, filename: str) -> Dict[str, float]:
+    url = "https://raw.githubusercontent.com/pytorch/test-infra/generated-stats/stats/test-times.json"
+
+    def process_response(the_response: Dict[str, Any]) -> Any:
+        build_environment = os.environ["BUILD_ENVIRONMENT"]
+        test_config = os.environ["TEST_CONFIG"]
+        return the_response[build_environment][test_config]
+
+    try:
+        return fetch_and_cache(dirpath, filename, url, process_response)
+    except Exception:
+        print("Couldn't download test times...")
+        return {}
+
+
 def get_disabled_tests(
     dirpath: str, filename: str = DISABLED_TESTS_FILE
 ) -> Optional[Dict[str, Any]]: