chore: Migrate gsutil usage to gcloud storage

Author: gurusai-voleti

benchmarks/maxtext_xpk_runner.py

@@ -428,7 +428,7 @@ def build_user_command(
   if wl_config.hlo_dump:
     hlo_dump = "XLA_FLAGS='--xla_dump_large_constants --xla_dump_to=/tmp/xla_dump'"
     upload_hlo_dump = (
-        f" && gsutil -m cp -r /tmp/xla_dump  {wl_config.base_output_directory}/{wl_config.run_name}/hlo_dump"
+        f" && gcloud storage cp --recursive /tmp/xla_dump  {wl_config.base_output_directory}/{wl_config.run_name}/hlo_dump"
     )
   # Construct the command string with proper formatting and line continuations
   command = " ".join(

benchmarks/upload_metrics_to_bq.py

@@ -187,7 +187,7 @@ def add_parser_arguments(parser: argparse.ArgumentParser):
 
 
 def download_metrics_file_locally(metrics_gcs_file: str, local_file: str) -> int:
-  command = f"gsutil cp -r {metrics_gcs_file} {local_file}"
+  command = f"gcloud storage cp --recursive {metrics_gcs_file} {local_file}"
   return run_command_with_updates(command, f"Download {metrics_gcs_file} in {local_file}")
 
 

src/dependencies/scripts/setup.sh

@@ -187,7 +187,7 @@ install_custom_libtpu() {
     # Install 'crcmod' to download 'libtpu.so' from GCS reliably
     python3 -m uv pip install -U crcmod
     # Copy libtpu.so from GCS path
-    gsutil cp "$LIBTPU_GCS_PATH" "$libtpu_path"
+    gcloud storage cp "$LIBTPU_GCS_PATH" "$libtpu_path"
 }
 
 install_maxtext_package_without_deps() {

src/maxtext/inference/mlperf/README.md

@@ -42,10 +42,10 @@ cd ${DATA_DISK_DIR?}
 #### LLama2-70b:
 
 ```
-gsutil cp gs://cloud-tpu-inference-public/mlcommons/inference/language/llama2-70b/data/processed-openorca/open_orca_gpt4_tokenized_llama.calibration_1000.pkl .
+gcloud storage cp gs://cloud-tpu-inference-public/mlcommons/inference/language/llama2-70b/data/processed-openorca/open_orca_gpt4_tokenized_llama.calibration_1000.pkl .
 mv open_orca_gpt4_tokenized_llama.calibration_1000.pkl processed-calibration-data.pkl
 
-gsutil cp gs://cloud-tpu-inference-public/mlcommons/inference/language/llama2-70b/data/processed-openorca/open_orca_gpt4_tokenized_llama.sampled_24576.pkl .
+gcloud storage cp gs://cloud-tpu-inference-public/mlcommons/inference/language/llama2-70b/data/processed-openorca/open_orca_gpt4_tokenized_llama.sampled_24576.pkl .
 mv open_orca_gpt4_tokenized_llama.sampled_24576.pkl processed-data.pkl
 ```
 

tests/post_training/integration/grpo_trainer_correctness_test.py

@@ -114,9 +114,10 @@ def setUp(self):
     super().setUp()
     jax.config.update("jax_default_prng_impl", "unsafe_rbg")
     command = [
-        "gsutil",
+        "gcloud",
+        "storage",
         "cp",
-        "-r",
+        "--recursive",
         "gs://maxtext-dataset/hf/llama3.1-tokenizer",
         os.path.join(MAXTEXT_ASSETS_ROOT, ""),
     ]

tests/post_training/unit/sft_data_processing_test.py

@@ -314,7 +314,7 @@
         {"test_data": QWEN_DATA},
     ]
 )
-@pytest.mark.external_training  # Uses gsutil to pull tokenizer.
+@pytest.mark.external_training  # Uses gcloud storage to pull tokenizer.
 class SFTDataProcessingTest(unittest.TestCase):
   test_data = {}
 
@@ -323,15 +323,16 @@ def setUpClass(cls):
     super().setUpClass()
     exit_code = subprocess.call(
         [
-            "gsutil",
+            "gcloud",
+            "storage",
             "cp",
-            "-r",
+            "--recursive",
             "gs://maxtext-dataset/hf/llama2-chat-tokenizer",
             os.path.join(MAXTEXT_ASSETS_ROOT, ""),
         ]
     )
     if exit_code != 0:
-      raise ValueError(f"Download tokenizer with gsutil cp failed with exit code: {exit_code}")
+      raise ValueError(f"Download tokenizer with gcloud storage cp failed with exit code: {exit_code}")
 
   def setUp(self):
     super().setUp()

tools/orchestration/multihost_job.py

@@ -104,7 +104,7 @@ def move_script_dir_to_gcs(script_dir, tmp_dir, zip_name, bucket_path):
 
   # Move zip file to GCS
   zip_in_gcs_path = "/".join((bucket_path, zip_name))
-  command = ["gsutil", "mv", zip_path, zip_in_gcs_path]
+  command = ["gcloud", "storage", "mv", zip_path, zip_in_gcs_path]
   captured_output = subprocess.run(command, check=True, capture_output=True)
 
   # Cleanup
@@ -148,7 +148,7 @@ def write_startup_script(zip_gcs_path, zip_name, log_name, bucket_path, startup_
 tar xzf {zip_name}
 {args.COMMAND}) 2>&1) >> {log_name}
 (echo "{finish_status_str()}") >> {log_name}
-gsutil cp {log_name} "{bucket_path}/"
+gcloud storage cp {log_name} "{bucket_path}/"
 (({create_kill_command_str(args)}) 2>&1 ) >> {log_name}"""
 
   with open(startup_script_file, "wt", encoding="utf-8") as f:
@@ -199,10 +199,10 @@ def write_download_from_gcs_sh(zip_gcs_path):
 while [ \$GCS_READ_SUCCESS -eq 0 ]
 do
   {{ # try
-      gsutil cp {zip_gcs_path} . &&
+      gcloud storage cp {zip_gcs_path} . &&
       echo 'Code download from GCS successful!' && GCS_READ_SUCCESS=1
   }} || {{ # catch
-      echo 'Failed to read GCS via gsutil, trying again'
+      echo 'Failed to read GCS via gcloud storage, trying again'
       sleep 10
   }}
 done"""
@@ -341,7 +341,7 @@ def main(raw_args=None) -> None:
   captured_output = move_script_dir_to_gcs(args.SCRIPT_DIR, tmp_dir_relative_to_script, zip_name, bucket_path)
   if captured_output.returncode != 0:
     print("\n\n Moving code to GCS failed")
-    print(f"Running 'gsutil mv zip {bucket_path}' failed with error: ")
+    print(f"Running 'gcloud storage mv zip {bucket_path}' failed with error: ")
     print(captured_output.stderr.decode())
     print("\nYou may need to run 'gcloud auth login'")
     return -1