Repository: kubeflow/pipelines

Issue:
[feature] Add helper function to change pipeline/operation ksa
### Feature Area

/area sdk

### What feature would you like to see?

I would like to see a helper function in the sdk to run pods using a kubernetes service account other than `pipeline-runner`. Discussed in slack at https://kubeflow.slack.com/archives/CE10KS9M4/p1629610669092400.

### What is the use case or pain point?

I'm using gke workload identity to map gcp service accounts to kubernetes service accounts, and I'm hoping to simplify the process of running different kfp workloads with different gcp iam access.

### Is there a workaround currently?

@Bobgy mentioned that this is possible in the generated sdk, but I'm not sure how that works yet.

<!-- Don't delete message below to encourage users to support your feature request! -->
Love this idea? Give it a 👍. We prioritize fulfilling features with the most 👍.

[feature] Add helper function to change pipeline/operation ksa
### Feature Area

/area sdk

### What feature would you like to see?

I would like to see a helper function in the sdk to run pods using a kubernetes service account other than `pipeline-runner`. Discussed in slack at https://kubeflow.slack.com/archives/CE10KS9M4/p1629610669092400.

### What is the use case or pain point?

I'm using gke workload identity to map gcp service accounts to kubernetes service accounts, and I'm hoping to simplify the process of running different kfp workloads with different gcp iam access.

### Is there a workaround currently?

@Bobgy mentioned that this is possible in the generated sdk, but I'm not sure how that works yet.

<!-- Don't delete message below to encourage users to support your feature request! -->
Love this idea? Give it a 👍. We prioritize fulfilling features with the most 👍.



Discussion:
Hi @jmcarp, thank you for raising the issue.
I can point you to related code in case you want to contribute, I think this can be a good first issue : ).

* KFP SDK client code: https://github.com/kubeflow/pipelines/blob/27b3511ad5e8283096555efb4f00d5ba5427af11/sdk/python/kfp/_client.py#L828
* How KFP SDK client uses the generated SDK API: https://github.com/kubeflow/pipelines/blob/27b3511ad5e8283096555efb4f00d5ba5427af11/sdk/python/kfp/_client.py#L659-L662
* Generated SDK API reference: https://kubeflow-pipelines.readthedocs.io/en/stable/source/kfp.server_api.html#kfp_server_api.api.run_service_api.RunServiceApi.create_run (Note the service_account field exists in `ApiRun` model.
* Generated SDK API tutorial: https://www.kubeflow.org/docs/components/pipelines/tutorials/sdk-examples/
@Bobgy I wrote up #6455, but it looks like we're thinking about different parts of the code. Does my PR make sense, or should I try a different approach?
> @Bobgy I wrote up #6455, but it looks like we're thinking about different parts of the code. Does my PR make sense, or should I try a different approach?

I tend to think service account as a submission time configuration. In that sense, adding the option to the client is probably better. That's also what we did for Vertex Pipelines client:
https://github.com/kubeflow/pipelines/blob/9afe125a759d8291f8a7732c6f81ca149c8c4ed5/sdk/python/kfp/v2/google/client/client.py#L272

#6455 would not be compatible with KFP v2 or Vertex Pipelines.
Thanks! I rewrote this as suggested--hope this makes more sense.
Hi @jmcarp, thank you for raising the issue.
I can point you to related code in case you want to contribute, I think this can be a good first issue : ).

* KFP SDK client code: https://github.com/kubeflow/pipelines/blob/27b3511ad5e8283096555efb4f00d5ba5427af11/sdk/python/kfp/_client.py#L828
* How KFP SDK client uses the generated SDK API: https://github.com/kubeflow/pipelines/blob/27b3511ad5e8283096555efb4f00d5ba5427af11/sdk/python/kfp/_client.py#L659-L662
* Generated SDK API reference: https://kubeflow-pipelines.readthedocs.io/en/stable/source/kfp.server_api.html#kfp_server_api.api.run_service_api.RunServiceApi.create_run (Note the service_account field exists in `ApiRun` model.
* Generated SDK API tutorial: https://www.kubeflow.org/docs/components/pipelines/tutorials/sdk-examples/
@Bobgy I wrote up #6455, but it looks like we're thinking about different parts of the code. Does my PR make sense, or should I try a different approach?
> @Bobgy I wrote up #6455, but it looks like we're thinking about different parts of the code. Does my PR make sense, or should I try a different approach?

I tend to think service account as a submission time configuration. In that sense, adding the option to the client is probably better. That's also what we did for Vertex Pipelines client:
https://github.com/kubeflow/pipelines/blob/9afe125a759d8291f8a7732c6f81ca149c8c4ed5/sdk/python/kfp/v2/google/client/client.py#L272

#6455 would not be compatible with KFP v2 or Vertex Pipelines.
Thanks! I rewrote this as suggested--hope this makes more sense.

Resolving patch:
diff --git a/sdk/python/kfp/_client.py b/sdk/python/kfp/_client.py
--- a/sdk/python/kfp/_client.py
+++ b/sdk/python/kfp/_client.py
@@ -692,6 +692,7 @@ def run_pipeline(
         version_id: Optional[str] = None,
         pipeline_root: Optional[str] = None,
         enable_caching: Optional[str] = None,
+        service_account: Optional[str] = None,
     ):
         """Run a specified pipeline.
 
@@ -715,6 +716,8 @@ def run_pipeline(
             individual tasks.
             If set, the setting applies to all tasks in the pipeline -- overrides
             the compile time settings.
+          service_account: Optional. Specifies which Kubernetes service account this
+            run uses.
 
         Returns:
           A run object. Most important field is id.
@@ -736,7 +739,8 @@ def run_pipeline(
         run_body = kfp_server_api.models.ApiRun(
             pipeline_spec=job_config.spec,
             resource_references=job_config.resource_references,
-            name=job_name)
+            name=job_name,
+            service_account=service_account)
 
         response = self._run_api.create_run(body=run_body)
 
@@ -765,6 +769,7 @@ def create_recurring_run(
         version_id: Optional[str] = None,
         enabled: bool = True,
         enable_caching: Optional[bool] = None,
+        service_account: Optional[str] = None,
     ):
         """Create a recurring run.
 
@@ -799,6 +804,8 @@ def create_recurring_run(
             individual tasks.
             If set, the setting applies to all tasks in the pipeline -- overrides
             the compile time settings.
+          service_account: Optional. Specifies which Kubernetes service account this
+            recurring run uses.
 
         Returns:
           A Job object. Most important field is id.
@@ -837,7 +844,8 @@ def create_recurring_run(
             description=description,
             no_catchup=no_catchup,
             trigger=trigger,
-            max_concurrency=max_concurrency)
+            max_concurrency=max_concurrency,
+            service_account=service_account)
         return self._job_api.create_job(body=job_body)
 
     def _create_job_config(
@@ -928,6 +936,7 @@ def create_run_from_pipeline_func(
         launcher_image: Optional[str] = None,
         pipeline_root: Optional[str] = None,
         enable_caching: Optional[bool] = None,
+        service_account: Optional[str] = None,
     ):
         """Runs pipeline on KFP-enabled Kubernetes cluster.
 
@@ -959,6 +968,8 @@ def create_run_from_pipeline_func(
             individual tasks.
             If set, the setting applies to all tasks in the pipeline -- overrides
             the compile time settings.
+          service_account: Optional. Specifies which Kubernetes service account this
+            run uses.
         """
         if pipeline_root is not None and mode == dsl.PipelineExecutionMode.V1_LEGACY:
             raise ValueError('`pipeline_root` should not be used with '
@@ -984,6 +995,7 @@ def create_run_from_pipeline_func(
                 namespace=namespace,
                 pipeline_root=pipeline_root,
                 enable_caching=enable_caching,
+                service_account=service_account,
             )
 
     def create_run_from_pipeline_package(
@@ -995,6 +1007,7 @@ def create_run_from_pipeline_package(
         namespace: Optional[str] = None,
         pipeline_root: Optional[str] = None,
         enable_caching: Optional[bool] = None,
+        service_account: Optional[str] = None,
     ):
         """Runs pipeline on KFP-enabled Kubernetes cluster.
 
@@ -1020,6 +1033,8 @@ def create_run_from_pipeline_package(
             individual tasks.
             If set, the setting applies to all tasks in the pipeline -- overrides
             the compile time settings.
+          service_account: Optional. Specifies which Kubernetes service account this
+            run uses.
         """
 
         class RunPipelineResult:
@@ -1060,6 +1075,7 @@ def __repr__(self):
             params=arguments,
             pipeline_root=pipeline_root,
             enable_caching=enable_caching,
+            service_account=service_account,
         )
         return RunPipelineResult(self, run_info)
 


Test patch: