[CI] Support using blob prefix for downloading So that uploading and downloading use the same file paths. Also refactor everything so that _get_blob_prefix is the common implementation.
diff --git a/.ci/cache_lit_timing_files.py b/.ci/cache_lit_timing_files.py index b44b826..ca223c2 100644 --- a/.ci/cache_lit_timing_files.py +++ b/.ci/cache_lit_timing_files.py
@@ -22,10 +22,13 @@ GCS_PARALLELISM = 100 +def _get_blob_prefix(): + return f"lit_timing_{platform.system().lower()}" + + def _maybe_upload_timing_file(bucket, timing_file_path): - blob_prefix = f"lit_timing_{platform.system().lower()}/" if os.path.exists(timing_file_path): - timing_file_blob = bucket.blob(blob_prefix + timing_file_path) + timing_file_blob = bucket.blob(_get_blob_prefix() + "/" + timing_file_path) timing_file_blob.upload_from_filename(timing_file_path) @@ -45,14 +48,14 @@ def _maybe_download_timing_file(blob): - file_name = blob.name.removeprefix("lit_timing/") + file_name = blob.name.removeprefix(_get_blob_prefix() + "/") pathlib.Path(os.path.dirname(file_name)).mkdir(parents=True, exist_ok=True) blob.download_to_filename(file_name) def download_timing_files(storage_client, bucket_name: str): bucket = storage_client.bucket(bucket_name) - blobs = bucket.list_blobs(prefix="lit_timing") + blobs = bucket.list_blobs(prefix=_get_blob_prefix()) with multiprocessing.pool.ThreadPool(GCS_PARALLELISM) as thread_pool: futures = [] for timing_file_blob in blobs: