Add support for non-default service accounts (#42)

buildstockbatch/gcp/gcp.py

@@ -609,9 +609,8 @@ def start_batch_job(self, batch_info):
         )
         instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate(policy=policy)
         allocation_policy = batch_v1.AllocationPolicy(instances=[instances])
-        # TODO: Add option to set service account that runs the job?
-        # Otherwise uses the project's default compute engine service account.
-        # allocation_policy.service_account = batch_v1.ServiceAccount(email = '')
+        if service_account := gcp_cfg.get("service_account"):
+            allocation_policy.service_account = batch_v1.ServiceAccount(email=service_account)
 
         # Define the batch job
         job = batch_v1.Job()
@@ -836,6 +835,7 @@ def start_combine_results_job_on_cloud(self, results_dir, do_timeseries=True):
                     ],
                     timeout=f"{60 * 60 * 24}s",  # 24h
                     max_retries=0,
+                    service_account=self.cfg["gcp"].get("service_account"),
                 )
             )
         )

buildstockbatch/schemas/v0.3.yaml

@@ -26,6 +26,7 @@ gcp-spec:
   job_identifier: regex('^[a-z]([a-z0-9-]{0,46}[a-z0-9])?$', required=True)
   project: str(required=True)
   region: str(required=True)
+  service_account: str(required=False)
   artifact_registry: include('gcp-ar-spec', required=True)
   batch_array_size: num(min=1, max=10000, required=True)
   gcs: include('gcs-spec', required=True)

docs/project_defn.rst

@@ -274,6 +274,8 @@ on the `GCP Batch <https://cloud.google.com/batch>`_ service.
    differentiate it from other jobs.
 *  ``project``: The GCP Project ID in which the batch will be run and of the Artifact Registry
    (where Docker images are stored).
+*  ``service_account``: The service account email address to use when running jobs on GCP.
+    Default: the Compute Engine default service account of the GCP project.
 *  ``gcs``: Configuration for project data storage on GCP Cloud Storage.
 
     *  ``bucket``: The Cloud Storage bucket this project will use for simulation output and