Support for managed job submission to Google Cloud - #64
Open
gregorkrz wants to merge 3 commits into
Open
Conversation
Add a `gcloud` launch scheduler so `pimm submit` can queue training on Google Cloud Batch: it provisions a single A100 VM, runs the published pimm container image, writes artifacts to a gcsfuse-mounted gs:// bucket (EXP_ROOT is rewritten to the local mount so training code stays storage-agnostic), and tears the VM down. - pimm/launch/gcloud.py: Batch job builder + submit driver, gs:// URI parsing, optional GCS code staging, job-name sanitization. - launch/sites/gcloud.yaml: documented gcloud site config. - config/schema/utils/submit: wire "gcloud" through validation, the scheduler enum, and the submit entrypoint. - tests/unit/test_launch_gcloud.py: unit coverage for the builder. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add a "Managed Google Cloud Batch Submission" section (prerequisites, `pimm submit --site gcloud` example, and gs:///staging/WANDB notes), list the site under File Ownership, and note the new `gcloud` scheduler value. Drop a stale gcloud.yaml comment referencing personal submit scripts. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
finalize_config now falls back to a WANDB_API_KEY=... line in the repo .env when neither --run.wandb-api-key nor an env-block value is set, so a gcloud submit authenticates to W&B via EITHER the .env line or the flag (the flag still wins). The .env file itself is never staged to the VM; the submit host reads the single key and injects it into the rendered Batch job env. Update gcloud.yaml, gcloud.py, and launch/README.md to document both paths. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
youngsm
reviewed
Jul 28, 2026
| # - A GCS bucket for output, and (for A100) quota in the chosen region. | ||
| # | ||
| # Fill in every REPLACE_ME_* value below before submitting. | ||
|
|
Member
There was a problem hiding this comment.
If we can, please remove a lot of the comments in this code. We can assume the user knows roughly what they're doing, so we don't need to be so overly explanatory.
youngsm
reviewed
Jul 28, 2026
|
|
||
|
|
||
| def wandb_api_key_from_dotenv(root: Path = ROOT) -> str | None: | ||
| """Read only WANDB_API_KEY from the repo `.env`, ignoring all other vars. |
Member
There was a problem hiding this comment.
Can this not be pulled in from os.getenv("WANDB_API_KEY")?
This function does not cover when users put their wandb api key in somewhere other than .env, for example their ~.bashrc.
Member
There was a problem hiding this comment.
Oh I see you check for that first then look for the dotenv. I think it would be better to just implement complete dotenv loading on import of pimm than have this code like this. That way you just have one line,
wandb_key = run_cfg.get("wandb_api_key") or cfg.get("env", {}).get("WANDB_API_KEY")
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds support for
--site gcloudto let users submit jobs to thhe Google Cloud Platform. The command pushes the current code to a Google Cloud Storage bucket and executes it in thepimmcontainer, using the machine type specified in the site config.