add security context to task execution metadata

Signed-off-by: noahjax <[email protected]>

dev-requirements.in

@@ -1,5 +1,5 @@
 -e file:.#egg=flytekit
-git+https://github.com/flyteorg/flyte.git@master#subdirectory=flyteidl
+git+https://github.com/dominodatalab/flyte.git@noahjax.add-owner-reference-to-create-task#subdirectory=flyteidl
 
 coverage[toml]
 hypothesis

flytekit/models/security.py

@@ -88,12 +88,14 @@ class Identity(_common.FlyteIdlEntity):
     iam_role: Optional[str] = None
     k8s_service_account: Optional[str] = None
     oauth2_client: Optional[OAuth2Client] = None
+    execution_identity: Optional[str] = None
 
     def to_flyte_idl(self) -> _sec.Identity:
         return _sec.Identity(
             iam_role=self.iam_role if self.iam_role else None,
             k8s_service_account=self.k8s_service_account if self.k8s_service_account else None,
             oauth2_client=self.oauth2_client.to_flyte_idl() if self.oauth2_client else None,
+            execution_identity=self.execution_identity,
         )
 
     @classmethod
@@ -104,6 +106,7 @@ def from_flyte_idl(cls, pb2_object: _sec.Identity) -> "Identity":
             oauth2_client=OAuth2Client.from_flyte_idl(pb2_object.oauth2_client)
             if pb2_object.oauth2_client and pb2_object.oauth2_client.ByteSize()
             else None,
+            execution_identity=pb2_object.execution_identity,
         )
 
 

flytekit/models/task.py

@@ -528,6 +528,7 @@ def __init__(
         annotations,
         k8s_service_account,
         environment_variables,
+        security_context,
     ):
         """
         Runtime task execution metadata.
@@ -539,13 +540,15 @@ def __init__(
         :param dict[str, str] annotations: Annotations to use for the execution of this task.
         :param Text k8s_service_account: Service account to use for execution of this task.
         :param dict[str, str] environment_variables: Environment variables for this task.
+        :param flytekit.models.security.SecurityContext security_context: Security context for this task, including run as, secrets, and tokens
         """
         self._task_execution_id = task_execution_id
         self._namespace = namespace
         self._labels = labels
         self._annotations = annotations
         self._k8s_service_account = k8s_service_account
         self._environment_variables = environment_variables
+        self._security_context = security_context
 
     @property
     def task_execution_id(self):
@@ -571,6 +574,10 @@ def k8s_service_account(self):
     def environment_variables(self):
         return self._environment_variables
 
+    @property
+    def security_context(self):
+        return self._security_context
+
     def to_flyte_idl(self):
         """
         :rtype: flyteidl.admin.agent_pb2.TaskExecutionMetadata
@@ -584,6 +591,7 @@ def to_flyte_idl(self):
             environment_variables={k: v for k, v in self.environment_variables.items()}
             if self.labels is not None
             else None,
+            security_context=self.security_context.to_flyte_idl() if self.security_context else None,
         )
         return task_execution_metadata
 
@@ -604,6 +612,9 @@ def from_flyte_idl(cls, pb2_object):
             environment_variables={k: v for k, v in pb2_object.environment_variables.items()}
             if pb2_object.environment_variables is not None
             else None,
+            security_context=_sec.SecurityContext.from_flyte_idl(pb2_object.security_context)
+            if pb2_object.security_context
+            else None,
         )
 
 

pyproject.toml

@@ -21,7 +21,7 @@ dependencies = [
     "diskcache>=5.2.1",
     "docker>=4.0.0,<7.0.0",
     "docstring-parser>=0.9.0",
-    "flyteidl>=1.11.0b1",
+    # "flyteidl>=1.11.0b1",
     "fsspec>=2023.3.0",
     "gcsfs>=2023.3.0",
     "googleapis-common-protos>=1.57",

tests/flytekit/unit/extend/test_agent.py

@@ -45,6 +45,7 @@
     WorkflowExecutionIdentifier,
 )
 from flytekit.models.literals import LiteralMap
+from flytekit.models.security import Identity, SecurityContext
 from flytekit.models.task import TaskExecutionMetadata, TaskTemplate
 from flytekit.tools.translator import get_serializable
 
@@ -157,6 +158,7 @@ def simple_task(i: int):
     annotations={"annotation_key": "annotation_val"},
     k8s_service_account="k8s service account",
     environment_variables={"env_var_key": "env_var_val"},
+    security_context=SecurityContext(run_as=Identity(execution_identity="task executor")),
 )