From cb15a4ee2d4a0faac69870858bc4803c04fdb3ae Mon Sep 17 00:00:00 2001
From: Saurabh Garg <144727284+iamsgarg-ob@users.noreply.github.com>
Date: Mon, 29 Jul 2024 11:42:02 -0700
Subject: [PATCH] [argo] add terminationMessagePolicy (#50) (#1936)

- add a terminationMessagePolicy to FallbackToLogsOnError
---
 metaflow/plugins/argo/argo_workflows.py           | 1 +
 metaflow/plugins/kubernetes/kubernetes_job.py     | 1 +
 metaflow/plugins/kubernetes/kubernetes_jobsets.py | 1 +
 3 files changed, 3 insertions(+)

diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py
index adc922635b2..7fa7f4a53e7 100644
--- a/metaflow/plugins/argo/argo_workflows.py
+++ b/metaflow/plugins/argo/argo_workflows.py
@@ -2015,6 +2015,7 @@ def _container_templates(self):
                         kubernetes_sdk.V1Container(
                             name=self._sanitize(node.name),
                             command=cmds,
+                            termination_message_policy="FallbackToLogsOnError",
                             ports=[kubernetes_sdk.V1ContainerPort(container_port=port)]
                             if port
                             else None,
diff --git a/metaflow/plugins/kubernetes/kubernetes_job.py b/metaflow/plugins/kubernetes/kubernetes_job.py
index 02d6e21dbd7..02645daeb40 100644
--- a/metaflow/plugins/kubernetes/kubernetes_job.py
+++ b/metaflow/plugins/kubernetes/kubernetes_job.py
@@ -98,6 +98,7 @@ def create_job_spec(self):
                     containers=[
                         client.V1Container(
                             command=self._kwargs["command"],
+                            termination_message_policy="FallbackToLogsOnError",
                             ports=[]
                             if self._kwargs["port"] is None
                             else [
diff --git a/metaflow/plugins/kubernetes/kubernetes_jobsets.py b/metaflow/plugins/kubernetes/kubernetes_jobsets.py
index a6c4c964b1c..25ec6a220c8 100644
--- a/metaflow/plugins/kubernetes/kubernetes_jobsets.py
+++ b/metaflow/plugins/kubernetes/kubernetes_jobsets.py
@@ -586,6 +586,7 @@ def dump(self):
                                 containers=[
                                     client.V1Container(
                                         command=self._kwargs["command"],
+                                        termination_message_policy="FallbackToLogsOnError",
                                         ports=[]
                                         if self._kwargs["port"] is None
                                         else [