From 0ecec660df5c0a5cf955ca8ad404f018333bb1d3 Mon Sep 17 00:00:00 2001
From: Lakhan Jindam <48401688+lakhanjindam@users.noreply.github.com>
Date: Sat, 8 Apr 2023 05:59:49 +0530
Subject: [PATCH] Fix validation error for retry strategy (#558)

Fix validation error for retry strategy

Signed-off-by: lakhanjindam <lakhanj569@gmail.com>
---
 .../workflows/upstream/retry_script.md        | 55 +++++++++++++++++++
 .../upstream/retry-script.upstream.yaml       | 20 +++++++
 examples/workflows/upstream/retry-script.yaml | 24 ++++++++
 examples/workflows/upstream/retry_script.py   | 14 +++++
 src/hera/workflows/retry_strategy.py          |  4 +-
 5 files changed, 115 insertions(+), 2 deletions(-)
 create mode 100644 docs/examples/workflows/upstream/retry_script.md
 create mode 100644 examples/workflows/upstream/retry-script.upstream.yaml
 create mode 100644 examples/workflows/upstream/retry-script.yaml
 create mode 100644 examples/workflows/upstream/retry_script.py

diff --git a/docs/examples/workflows/upstream/retry_script.md b/docs/examples/workflows/upstream/retry_script.md
new file mode 100644
index 000000000..a31ef56d0
--- /dev/null
+++ b/docs/examples/workflows/upstream/retry_script.md
@@ -0,0 +1,55 @@
+# Retry Script
+
+> Note: This example is a replication of an Argo Workflow example in Hera. The upstream example can be [found here](https://github.com/argoproj/argo-workflows/blob/master/examples/retry-script.yaml).
+
+
+
+
+=== "Hera"
+
+    ```python linenums="1"
+    from hera.workflows import Workflow, script, RetryStrategy
+
+
+    @script(image="python:alpine3.6", retry_strategy=RetryStrategy(limit=10), add_cwd_to_sys_path=False)
+    def retry_script():
+        import random
+        import sys
+
+        exit_code = random.choice([0, 1, 1])
+        sys.exit(exit_code)
+
+
+    with Workflow(generate_name="retry-script-", entrypoint="retry-script") as w:
+        retry_script()
+    ```
+
+=== "YAML"
+
+    ```yaml linenums="1"
+    apiVersion: argoproj.io/v1alpha1
+    kind: Workflow
+    metadata:
+      generateName: retry-script-
+    spec:
+      entrypoint: retry-script
+      templates:
+      - name: retry-script
+        retryStrategy:
+          limit: '10'
+        script:
+          command:
+          - python
+          image: python:alpine3.6
+          source: 'import random
+
+            import sys
+
+
+            exit_code = random.choice([0, 1, 1])
+
+            sys.exit(exit_code)
+
+            '
+    ```
+
diff --git a/examples/workflows/upstream/retry-script.upstream.yaml b/examples/workflows/upstream/retry-script.upstream.yaml
new file mode 100644
index 000000000..92bc489d0
--- /dev/null
+++ b/examples/workflows/upstream/retry-script.upstream.yaml
@@ -0,0 +1,20 @@
+# This example demonstrates the use of retries for a single script.
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+metadata:
+  generateName: retry-script-
+spec:
+  entrypoint: retry-script
+  templates:
+  - name: retry-script
+    retryStrategy:
+      limit: "10"
+    script:
+      image: python:alpine3.6
+      command: ["python"]
+      # fail with a 66% probability
+      source: |
+        import random;
+        import sys; 
+        exit_code = random.choice([0, 1, 1]); 
+        sys.exit(exit_code)
diff --git a/examples/workflows/upstream/retry-script.yaml b/examples/workflows/upstream/retry-script.yaml
new file mode 100644
index 000000000..b62e1b099
--- /dev/null
+++ b/examples/workflows/upstream/retry-script.yaml
@@ -0,0 +1,24 @@
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+metadata:
+  generateName: retry-script-
+spec:
+  entrypoint: retry-script
+  templates:
+  - name: retry-script
+    retryStrategy:
+      limit: '10'
+    script:
+      command:
+      - python
+      image: python:alpine3.6
+      source: 'import random
+
+        import sys
+
+
+        exit_code = random.choice([0, 1, 1])
+
+        sys.exit(exit_code)
+
+        '
diff --git a/examples/workflows/upstream/retry_script.py b/examples/workflows/upstream/retry_script.py
new file mode 100644
index 000000000..9fb59a327
--- /dev/null
+++ b/examples/workflows/upstream/retry_script.py
@@ -0,0 +1,14 @@
+from hera.workflows import RetryStrategy, Workflow, script
+
+
+@script(image="python:alpine3.6", retry_strategy=RetryStrategy(limit=10), add_cwd_to_sys_path=False)
+def retry_script():
+    import random
+    import sys
+
+    exit_code = random.choice([0, 1, 1])
+    sys.exit(exit_code)
+
+
+with Workflow(generate_name="retry-script-", entrypoint="retry-script") as w:
+    retry_script()
diff --git a/src/hera/workflows/retry_strategy.py b/src/hera/workflows/retry_strategy.py
index f277e45ac..dc79bffea 100644
--- a/src/hera/workflows/retry_strategy.py
+++ b/src/hera/workflows/retry_strategy.py
@@ -36,7 +36,7 @@ class RetryStrategy(_BaseModel):
     affinity: Optional[RetryAffinity] = None
     backoff: Optional[Backoff] = None
     expression: Optional[str] = None
-    limit: Optional[Union[int, str]] = None
+    limit: Optional[Union[str, int, IntOrString]] = None
     retry_policy: Optional[Union[str, RetryPolicy]] = None
 
     @validator("retry_policy", pre=True)
@@ -52,7 +52,7 @@ def _convert_limit(cls, v):
         if v is None or isinstance(v, IntOrString):
             return v
 
-        return IntOrString(__root__=str(v))  # int or str
+        return str(v)  # int or str
 
     def build(self) -> _ModelRetryStrategy:
         return _ModelRetryStrategy(