Add addition attributes in RunSpec

src/flyte/_run.py

@@ -129,7 +129,7 @@ def __init__(
     async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
         import grpc
         from flyteidl2.common import identifier_pb2
-        from flyteidl2.core import literals_pb2
+        from flyteidl2.core import literals_pb2, security_pb2
         from flyteidl2.task import run_pb2
         from flyteidl2.workflow import run_definition_pb2, run_service_pb2
         from google.protobuf import wrappers_pb2
@@ -258,6 +258,14 @@ async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args,
             env_kv = run_pb2.Envs(values=kv_pairs)
             annotations = run_pb2.Annotations(values=self._annotations)
             labels = run_pb2.Labels(values=self._labels)
+            raw_data_storage = (
+                run_pb2.RawDataStorage(raw_data_prefix=self._raw_data_path) if self._raw_data_path else None
+            )
+            security_context = (
+                security_pb2.SecurityContext(run_as=security_pb2.Identity(k8s_service_account=self._service_account))
+                if self._service_account
+                else None
+            )
 
             try:
                 resp = await get_client().run_service.CreateRun(
@@ -275,6 +283,8 @@ async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args,
                             labels=labels,
                             envs=env_kv,
                             cluster=self._queue or task.queue,
+                            raw_data_storage=raw_data_storage,
+                            security_context=security_context,
                         ),
                     ),
                 )
@@ -603,8 +613,8 @@ async def example_task(x: int, y: str) -> str:
     :param interactive_mode: Optional, can be forced to True or False.
          If not provided, it will be set based on the current environment. For example Jupyter notebooks are considered
          interactive mode, while scripts are not. This is used to determine how the code bundle is created.
-    :param raw_data_path: Use this path to store the raw data for the run. Currently only supported for local runs,
-      and can be used to store raw data in specific locations. TODO coming soon for remote runs as well.
+    :param raw_data_path: Use this path to store the raw data for the run for local and remote, and can be used to
+         store raw data in specific locations.
     :param run_base_dir: Optional The base directory to use for the run. This is used to store the metadata for the run,
      that is passed between tasks.
     :param overwrite_cache: Optional If true, the cache will be overwritten for the run

src/flyte/cli/_run.py

@@ -86,6 +86,26 @@ class RunArguments:
             )
         },
     )
+    raw_data_path: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--raw-data-path"],
+                type=str,
+                help="Override the output prefix used to store offloaded data types. e.g. s3://bucket/",
+            )
+        },
+    )
+    service_account: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--service-account"],
+                type=str,
+                help="Kubernetes service account. If not provided, the configured default will be used",
+            )
+        },
+    )
     name: str | None = field(
         default=None,
         metadata={
@@ -140,6 +160,8 @@ async def _run():
                 copy_style=self.run_args.copy_style,
                 mode="local" if self.run_args.local else "remote",
                 name=self.run_args.name,
+                raw_data_path=self.run_args.raw_data_path,
+                service_account=self.run_args.service_account,
             ).run.aio(self.obj, **ctx.params)
             if self.run_args.local:
                 console.print(