Skip to content

Reverts f7f720404a0c6cd47133419edd1a8ff4c5effa3c #5174

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Reverts f7f720404a0c6cd47133419edd1a8ff4c5effa3c #5174

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 19 additions & 38 deletions src/maxtext/utils/gcs_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,12 +76,7 @@ def add_trailing_slash(path):


def upload_blob(destination_gcs_name, source_file_name):
"""Uploads a file to a GCS or filesystem/CNS location (no-op if not found and decoupled)."""
if not destination_gcs_name.startswith("gs://"):
dest = epath.Path(destination_gcs_name)
dest.parent.mkdir(parents=True, exist_ok=True)
epath.Path(source_file_name).copy(dest, overwrite=True)
return
"""Uploads a file to a GCS location (no-op if not found and decoupled)."""
if not _gcs_guard("upload_blob"):
return
bucket_name, prefix_name = parse_gcs_bucket_and_prefix(destination_gcs_name)
Expand All @@ -92,44 +87,30 @@ def upload_blob(destination_gcs_name, source_file_name):


def upload_dump(local_dir, target_dir, module_name=None, delete_local_after=True, all_host_upload=False):
"""Uploads a directory to a GCS or filesystem/CNS location, with an optional filter (no-op if not found and decoupled)."""
"""Uploads a directory to a GCS location, with an optional filter (no-op if not found and decoupled)."""
if not all_host_upload and jax.process_index() != 0:
return
if not _gcs_guard("upload_dump"):
return
storage_client = storage.Client()
bucket_name, prefix_name = parse_gcs_bucket_and_prefix(target_dir)
bucket = storage_client.get_bucket(bucket_name)
if all_host_upload:
hostname = socket.gethostname() # Alternatively can use jax.process_id()
prefix_name = os.path.join(prefix_name, hostname)
target_dir = os.path.join(target_dir, hostname)
max_logging.log(f"Uploading Dump to {target_dir}...")
if not target_dir.startswith("gs://"):
dest_dir = epath.Path(target_dir)
dest_dir.mkdir(parents=True, exist_ok=True)
for root, _, files in os.walk(local_dir):
for file in files:
if module_name and module_name not in file:
continue
else:
max_logging.log(f"Uploading {file}")
local_path = os.path.join(root, file)
relative_path = os.path.relpath(local_path, local_dir)
dest_path = dest_dir / relative_path
dest_path.parent.mkdir(parents=True, exist_ok=True)
epath.Path(local_path).copy(dest_path, overwrite=True)
else:
if not _gcs_guard("upload_dump"):
return
storage_client = storage.Client()
bucket_name, prefix_name = parse_gcs_bucket_and_prefix(target_dir)
bucket = storage_client.get_bucket(bucket_name)
for root, _, files in os.walk(local_dir):
for file in files:
if module_name and module_name not in file:
continue
else:
max_logging.log(f"Uploading {file}")
local_path = os.path.join(root, file)
relative_path = os.path.relpath(local_path, local_dir)
blob_name = os.path.join(prefix_name, relative_path)
blob = bucket.blob(blob_name)
blob.upload_from_filename(local_path)
for root, _, files in os.walk(local_dir):
for file in files:
if module_name and module_name not in file:
continue
else:
max_logging.log(f"Uploading {file}")
local_path = os.path.join(root, file)
relative_path = os.path.relpath(local_path, local_dir)
blob_name = os.path.join(prefix_name, relative_path)
blob = bucket.blob(blob_name)
blob.upload_from_filename(local_path)
max_logging.log(f"Dump Uploaded to {target_dir}!")
if delete_local_after:
shutil.rmtree(local_dir)
Expand Down
Loading