[batch] Use gcloud to localize large input files from GCS

[kush-chandra]

Change Description

Fixes #15011

This is a redo of #15273 which addresses the scalability issues with that approach:

• We now only use the chunked gcloud downloader if the file is large enough to benefit. Files below the chunk size will be downloaded in a single process, as before.
• Both the gcloud download paths are now bounded by thread count & memory buffer usage, preventing memory errors on lowmem VMs.

The timing improvements for large files remain the same as in the original change, and there is now no spike in latency when downloading hundreds of smaller files.

Security Assessment

• This change potentially impacts the Hail Batch instance as deployed by Broad Institute in GCP

Impact Rating

• This change has a medium security impact

Impact Description

This change instantiates a new Google client to download files from GCS buckets. It uses the same local credential files as the existing custom clients we've created. I manually verified the problem cases with the original change now work as intended and am working on adding additional testing for those paths.

Appsec Review

• Required: The impact has been assessed and approved by appsec

[cjllanwarne]

Can we check ahead of time that the directory exists (instead of letting it throw an exception that we then have to catch do resubmit the exact same command in the except block?)

[kush-chandra]

Yep, I've refactored this. The original copier code used that pattern to rely on an implementation detail of local_fs to catch an edge case where we're trying to expand a directory into an existing file which is not a directory, so I made that check explicit here.

[cjllanwarne]

Same thought here

[cjllanwarne]

There's two risky things here:

1. We already acquired a sema slot from bounded_gather2. Especially with the xfer_sema waits introducing an easy place for tasks to be paused/suspended, it's very possible that lots of tasks get a bounded_gather2 sema slots and so the sema pool is empty before we try to acquire another 8 slots for the download here
2. Because this acquiring logic is not atomic, it's possible that a thread might pick up a sub-portion of its 8 required slots here, then get suspended, and those reserved slots are locked up in this task which isn't doing anything with them yet. It looks like the weighted semaphor semantics is maybe more what you're looking for here (you can atomically request a number of slots)?

[kush-chandra]

I changed the copier semaphore to a weighted semaphore so this can atomically request 8 threads from the shared pool. xfer_sema appears to allow more simultaneous transfers than sema does, so I think this mitigates the deadlock risk. (The original copier has the same thresholds for sema & xfer_sema, so I'm wondering how much value xfer_sema is actually providing like this.)

[cjllanwarne]

I believe these access tokens have a 1 hour lifespan before they'll be rejected. We can split on the type of google credentials we have and use one of Credentials.from_service_account_info(credentials.key) for service accounts or Credentials(token=None, refresh_token=credentials.credentials['refresh_token'], client_id=credentials.credentials['client_id'], client_secret=credentials.credentials['client_secret'], token_uri='https://oauth2.googleapis.com/token') for application-default credentials, or Client(credentials = None) for automatic application-default credentials.

In fact, credentials.access_token looks like a method (and an async method at that), so this would return a method to produce a coroutine rather than a value, so I'm not sure this code path would work even temporarily. I think we must be going the credentials=None / application-default route below during the on-VM testing, but I'm not sure it'd work if we tried to pass credentials in on a local laptop

[kush-chandra]

Thanks for catching that. I split it up into service account/token creds/anonymous creds/application default creds. From what I could tell, VM testing is indeed using the default path.

[graphite-app]

Blocking I/O operations (os.path.exists and os.makedirs) are called in an async method without being wrapped in blocking_to_async. This blocks the event loop.

if not os.path.exists(os.path.dirname(dest)):
    os.makedirs(os.path.dirname(dest), exist_ok=True)
# Should be:
await blocking_to_async(self._thread_pool, os.makedirs, os.path.dirname(dest), exist_ok=True)

Also, the os.path.exists check is redundant since exist_ok=True handles existing directories.

Suggested change

	if not os.path.exists(os.path.dirname(dest)):
	os.makedirs(os.path.dirname(dest), exist_ok=True)
	await blocking_to_async(self._thread_pool, os.makedirs, os.path.dirname(dest), exist_ok=True)

Spotted by Graphite



Is this helpful? React 👍 or 👎 to let us know.

[cjllanwarne]

If #15395 merges, the test batch for this PR would only have run the following steps:

[
    'check_hail',
    'check_pip_requirements',
    'check_services',
    'merge_code',
    'test_auth',
    'test_auth_copy_paste_login',
    'test_auth_copy_paste_login_timeout',
    'test_batch',
    'test_batch_docs',
    'test_batch_invariants',
    'test_batch_job_private_machines',
    'test_ci',
    'test_ci_unit',
    'test_hail_python',
    'test_hail_python_local_backend',
    'test_hail_python_service_backend_gcp',
    'test_hail_python_unchecked_allocator',
    'test_hailctl_batch',
    'test_hailtop_python',
    'test_hailtop_python_fs',
    'test_monitoring',
]

[cjllanwarne]

I have a few thoughts -

1. I don't fully understand what's happening here. Would you be willing to give a walkthrough?
2. I'm a little scared given how many big - and pretty foundational - changes are happening to such an intricate (brittle!) system
3. One of my original big hopes was to reduce the maintenance overhead of our in-house copier having to re-implement gcloud storage cp in python code. But we are doubling down on that in this PR

Before we commit to this route... I'd really like to see what a "shell out to gcloud storage cp" version would look like. Even if we have to make a couple of calls because it doesn't support many <=> many directory transfers, that still feels like it would be better than continuing to follow this "re-implement it ourselves" codebase. What do you think?