keep container validation

Signed-off-by: byhsu <[email protected]>
flyteorg · Feb 1, 2023 · d3b793b · d3b793b
1 parent 60164d6
commit d3b793b
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 15 deletions.
diff --git a/flytekit/core/pod_template.py b/flytekit/core/pod_template.py
@@ -8,19 +8,9 @@
 PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
 
 
-# overwrite setter to bypass unnecessary validation
-# https://github.com/kubernetes-client/python/blob/master/kubernetes/client/models/v1_pod_spec.py#L323
-@V1PodSpec.containers.setter
-def custom_containers_setter(self, containers):
-    self._containers = containers
-
-
-V1PodSpec.containers = custom_containers_setter
-
-
 @dataclass
 class PodTemplate(object):
-    pod_spec: V1PodSpec = V1PodSpec()
+    pod_spec: V1PodSpec = V1PodSpec(containers=[])
     primary_container_name: str = PRIMARY_CONTAINER_DEFAULT_NAME
     labels: Optional[Dict[str, str]] = None
     annotations: Optional[Dict[str, str]] = None

diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py
@@ -209,7 +209,7 @@ def _get_container(self, settings: SerializationSettings) -> _task_model.Contain
         )
 
     def _serialize_pod_spec(self, settings: SerializationSettings) -> Dict[str, Any]:
-        containers = self.pod_template.pod_spec.containers if self.pod_template.pod_spec.containers is not None else []
+        containers = self.pod_template.pod_spec.containers
         primary_exists = False
 
         for container in containers:

diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py
@@ -233,9 +233,7 @@ def test_pod_template(default_serialization_settings):
         primary_container_name="primary",
         labels={"lKeyA": "lValA"},
         annotations={"aKeyA": "aValA"},
-        pod_spec=V1PodSpec(
-            volumes=[V1Volume(name="volume")],
-        ),
+        pod_spec=V1PodSpec(volumes=[V1Volume(name="volume")], containers=[]),
     ),
     pod_template_name="A",
 )