feat(job): implement custom environment variables (--env flag) support

[Neelabh94]

• New --env Flag: Added a new --env flag to the 'gcluster job submit' command to allow users to pass custom environment variables to their workload containers.

• Validation Logic: Implemented validation for environment variable keys to ensure they adhere to POSIX naming conventions (letters, numbers, and underscores, without starting with a digit).

• Template Updates: Updated GKE and Pathways JobSet templates to inject the provided environment variables into the container specifications.

• Testing and Documentation: Added comprehensive unit tests for flag parsing and validation, and updated the job guide documentation with usage examples.

Submission Checklist

NOTE: Community submissions can take up to 2 weeks to be reviewed.

Please take the following actions before submitting this pull request.

• Fork your PR branch from the Toolkit "develop" branch (not main)
• Test all changes with pre-commit in a local branch #
• Confirm that "make tests" passes all tests
• Add or modify unit tests to cover code changes
• Ensure that unit test coverage remains above 80%
• Update all applicable documentation
• Follow Cluster Toolkit Contribution guidelines #

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enables greater flexibility for job submissions by allowing users to define custom environment variables via the --env flag. The changes include necessary updates to the CLI command, orchestration logic, and Kubernetes manifest templates, ensuring that environment variables are correctly propagated to the target containers while maintaining strict naming validation.

Highlights

• New --env Flag: Added a new --env flag to the 'gcluster job submit' command to allow users to pass custom environment variables to their workload containers.
• Validation Logic: Implemented validation for environment variable keys to ensure they adhere to POSIX naming conventions (letters, numbers, and underscores, without starting with a digit).
• Template Updates: Updated GKE and Pathways JobSet templates to inject the provided environment variables into the container specifications.
• Testing and Documentation: Added comprehensive unit tests for flag parsing and validation, and updated the job guide documentation with usage examples.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for passing custom environment variables to workload containers via a new --env flag in the gcluster job submit command. It includes POSIX-compliant validation for environment variable keys, updates the Kubernetes JobSet and Pathways templates to inject these variables, and adds comprehensive unit tests and documentation. The reviewer's feedback focuses on ensuring deterministic manifest generation by converting the environment variables map into a sorted slice of structs before rendering the templates, preventing randomized ordering in the generated Kubernetes manifests.

[Neelabh94]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for custom environment variables in the 'gcluster job submit' command via a new '--env' flag. The changes include command-line flag parsing and validation, updates to GKE and Pathways JobSet templates to inject the environment variables, comprehensive unit tests, and updated documentation. The review feedback recommends minor performance optimizations to avoid unnecessary map and slice allocations by returning early when no environment variables are provided.

[Neelabh94]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for passing custom environment variables to workload containers via a new --env flag in the gcluster job submit command. It includes input validation to ensure environment variable keys conform to POSIX naming standards, updates the GKE and Pathways JobSet templates to inject these variables, and adds comprehensive unit tests and documentation. I have no feedback to provide as there are no review comments to evaluate.

[Neelabh94]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for passing custom environment variables to workload containers via a new --env flag in the gcluster job submit command, complete with POSIX-compliant validation, template updates for GKE and Pathways, unit tests, and documentation. The review feedback correctly points out that for Pathways jobs, these environment variables must also be propagated to the pathways-worker container in the worker replicatedJob to ensure critical JAX/XLA configurations are not ignored during TPU execution.

[gemini-code-assist]

While injecting custom environment variables into the workload-container on the head node is necessary, for Pathways jobs, the actual JAX/TPU computations are executed by the pathways-worker containers in the worker replicatedJob (defined further down in this template).

Many critical JAX, XLA, and PJRT configuration environment variables (such as XLA_FLAGS, JAX_XLA_FLAGS, or debug flags) must be present on the workers to take effect. If they are only set on the head node's workload-container, they will be completely ignored during TPU execution.

Please also propagate $.Env to the pathways-worker container's env block in the worker replicatedJob. Note that if you do this, you will also need to update the test TestGeneratePathwaysManifest_CustomEnv in pkg/orchestrator/gke/gke_job_orchestrator_test.go to expect the custom environment variables to appear in both containers.