Skip to content

Add network tags support for Google Batch #5951

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.

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
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 5 additions & 0 deletions docs/reference/config.md
Original file line number Diff line number Diff line change
Expand Up @@ -882,6 +882,11 @@ The following settings are available for Google Cloud Batch:
- projects/{project}/global/networks/{network}
- global/networks/{network}

`google.batch.networkTags`
: The network tags to be applied to the instances created by Google Batch jobs. Network tags are used to apply firewall rules and control network access (e.g., `['allow-ssh', 'allow-http']`).

: Network tags are ignored when using instance templates. See [Add network tags](https://cloud.google.com/vpc/docs/add-remove-network-tags) for more information.

`google.batch.serviceAccountEmail`
: Define the Google service account email to use for the pipeline execution. If not specified, the default Compute Engine service account for the project will be used.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,14 @@ Max number of execution attempts of a job interrupted by a Compute Engine spot r
""")
public String network;

@ConfigOption
@Description("""
The network tags to be applied to the instances created by Google Batch jobs (e.g., `['allow-ssh', 'allow-http']`).

[Read more](https://cloud.google.com/vpc/docs/add-remove-network-tags)
""")
public List<String> networkTags;

@ConfigOption
@Description("""
The Google service account email to use for the pipeline execution. If not specified, the default Compute Engine service account for the project will be used.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,10 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {

allocationPolicy.putAllLabels( task.config.getResourceLabels() )

// Add network tags if configured
if( executor.config.networkTags )
allocationPolicy.addAllTags( executor.config.networkTags )

// use instance template if specified
if( task.config.getMachineType()?.startsWith('template://') ) {
if( task.config.getAccelerator() )
Expand All @@ -307,6 +311,9 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
if( executor.config.cpuPlatform )
log.warn1 'Config option `google.batch.cpuPlatform` ignored because an instance template was specified'

if( executor.config.networkTags )
log.warn1 'Config option `google.batch.networkTags` ignored because an instance template was specified'

if( executor.config.preemptible )
log.warn1 'Config option `google.batch.premptible` ignored because an instance template was specified'

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ class BatchConfig {
private String network
private String subnetwork
private String serviceAccountEmail
private List<String> networkTags
private BatchRetryConfig retryConfig
private List<Integer> autoRetryExitCodes
private List<String> gcsfuseOptions
Expand All @@ -69,6 +70,7 @@ class BatchConfig {
String getNetwork() { network }
String getSubnetwork() { subnetwork }
String getServiceAccountEmail() { serviceAccountEmail }
List<String> getNetworkTags() { networkTags }
BatchRetryConfig getRetryConfig() { retryConfig }
List<Integer> getAutoRetryExitCodes() { autoRetryExitCodes }
List<String> getGcsfuseOptions() { gcsfuseOptions }
Expand All @@ -89,6 +91,7 @@ class BatchConfig {
result.network = session.config.navigate('google.batch.network')
result.subnetwork = session.config.navigate('google.batch.subnetwork')
result.serviceAccountEmail = session.config.navigate('google.batch.serviceAccountEmail')
result.networkTags = session.config.navigate('google.batch.networkTags', List.of()) as List<String>
result.retryConfig = new BatchRetryConfig( session.config.navigate('google.batch.retryPolicy') as Map ?: Map.of() )
result.autoRetryExitCodes = session.config.navigate('google.batch.autoRetryExitCodes', DEFAULT_RETRY_LIST) as List<Integer>
result.gcsfuseOptions = session.config.navigate('google.batch.gcsfuseOptions', DEFAULT_GCSFUSE_OPTS) as List<String>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,7 @@ class GoogleBatchTaskHandlerTest extends Specification {
getAutoRetryExitCodes() >> [50001,50002]
getSpot() >> true
getNetwork() >> 'net-1'
getNetworkTags() >> ['tag1', 'tag2']
getServiceAccountEmail() >> 'foo@bar.baz'
getSubnetwork() >> 'subnet-1'
getUsePrivateAddress() >> true
Expand Down Expand Up @@ -219,6 +220,7 @@ class GoogleBatchTaskHandlerTest extends Specification {
allocationPolicy.getInstances(0).getInstallGpuDrivers() == true
allocationPolicy.getLabelsMap() == [foo: 'bar']
allocationPolicy.getServiceAccount().getEmail() == 'foo@bar.baz'
allocationPolicy.getTagsList() == ['tag1', 'tag2']
and:
instancePolicy.getAccelerators(0).getCount() == 1
instancePolicy.getAccelerators(0).getType() == ACCELERATOR.type
Expand Down
Loading