.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto/core/control_flow/conditions.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_core_control_flow_conditions.py: Conditions -------------- Flytekit supports conditions as a first class construct in the language. Conditions offer a way to selectively execute branches of a workflow based on static or dynamic data produced by other tasks or come in as workflow inputs. Conditions are very performant to be evaluated. However, they are limited to certain binary and logical operators and can only be performed on primitive values. .. GENERATED FROM PYTHON SOURCE LINES 11-12 Import the necessary modules. .. GENERATED FROM PYTHON SOURCE LINES 12-17 .. code-block:: default import random from flytekit import conditional, task, workflow .. GENERATED FROM PYTHON SOURCE LINES 18-22 Example 1 ^^^^^^^^^ In this example, we define two tasks `square` and `double`. Depending on whether the workflow input is a fraction (0-1) or not, the respective task is executed. .. GENERATED FROM PYTHON SOURCE LINES 22-62 .. code-block:: default @task def square(n: float) -> float: """ Parameters: n (float): name of the parameter for the task is derived from the name of the input variable, and the type is automatically mapped to Types.Integer Return: float: The label for the output is automatically assigned and the type is deduced from the annotation """ return n * n @task def double(n: float) -> float: """ Parameters: n (float): name of the parameter for the task is derived from the name of the input variable and the type is mapped to ``Types.Integer`` Return: float: The label for the output is auto-assigned and the type is deduced from the annotation """ return 2 * n @workflow def multiplier(my_input: float) -> float: return ( conditional("fractions") .if_((my_input >= 0.1) & (my_input <= 1.0)) .then(double(n=my_input)) .else_() .then(square(n=my_input)) ) if __name__ == "__main__": print(f"Output of multiplier(my_input=3.0): {multiplier(my_input=3.0)}") print(f"Output of multiplier(my_input=0.5): {multiplier(my_input=0.5)}") .. GENERATED FROM PYTHON SOURCE LINES 63-72 Example 2 ^^^^^^^^^ In this example, we define an ``if`` condition with multiple branches. It fails if none of the conditions is met. Flyte expects any ``conditional()`` statement to be **complete**. This means all possible branches should be handled. .. note:: Notice the use of bitwise (&). Python (PEP-335) doesn't allow overloading of the logical ``and``, ``or``, and ``not`` operators. Flytekit uses bitwise `&` and `|` as logical ``and`` and ``or`` operators. This is a common practice in other libraries too. .. GENERATED FROM PYTHON SOURCE LINES 72-89 .. code-block:: default @workflow def multiplier_2(my_input: float) -> float: return ( conditional("fractions") .if_((my_input > 0.1) & (my_input < 1.0)) .then(double(n=my_input)) .elif_((my_input > 1.0) & (my_input <= 10.0)) .then(square(n=my_input)) .else_() .fail("The input must be between 0 and 10") ) if __name__ == "__main__": print(f"Output of multiplier_2(my_input=10.0): {multiplier_2(my_input=10.0)}") .. GENERATED FROM PYTHON SOURCE LINES 90-93 Example 3 ^^^^^^^^^ In this example, we consume the output returned by the ``conditional()`` in the subsequent task. .. GENERATED FROM PYTHON SOURCE LINES 93-114 .. code-block:: default @workflow def multiplier_3(my_input: float) -> float: result = ( conditional("fractions") .if_((my_input > 0.1) & (my_input < 1.0)) .then(double(n=my_input)) .elif_((my_input > 1.0) & (my_input < 10.0)) .then(square(n=my_input)) .else_() .fail("The input must be between 0 and 10") ) # the 'result' will either be the output of `double` or `square`. If none of the conditions is true, # it gives a failure message. return double(n=result) if __name__ == "__main__": print(f"Output of multiplier_3(my_input=5.0): {multiplier_3(my_input=5.0)}") .. GENERATED FROM PYTHON SOURCE LINES 115-126 Example 4 ^^^^^^^^^^ It is possible to test if a boolean returned from the previous task is True. But unary operations are not supported. Use the `is_true`, `is_false` or `is_` on the result instead. .. note:: How do output values get these methods? In a workflow, no output can be accessed directly. The inputs and outputs are auto-wrapped in a special object called :py:class:`flytekit.extend.Promise`. In this example, we create a biased coin whose seed can be controlled. .. GENERATED FROM PYTHON SOURCE LINES 126-161 .. code-block:: default @task def coin_toss(seed: int) -> bool: """ Mimic some condition to check if the operation was successfully executed. """ r = random.Random(seed) if r.random() < 0.5: return True return False @task def failed() -> int: """ Mimic a task that handles failure """ return -1 @task def success() -> int: """ Mimic a task that handles success """ return 0 @workflow def basic_boolean_wf(seed: int = 5) -> int: result = coin_toss(seed=seed) return ( conditional("test").if_(result.is_true()).then(success()).else_().then(failed()) ) .. GENERATED FROM PYTHON SOURCE LINES 162-171 Example 5 ^^^^^^^^^^ It is possible to pass a boolean directly to a workflow. .. note:: Note that the boolean passed has a method named `is_true`. This boolean is present within the workflow context and is wrapped in a Flytekit special object. This special object allows it to have the additional behavior. .. GENERATED FROM PYTHON SOURCE LINES 171-176 .. code-block:: default @workflow def bool_input_wf(b: bool) -> int: return conditional("test").if_(b.is_true()).then(success()).else_().then(failed()) .. GENERATED FROM PYTHON SOURCE LINES 177-178 The workflow can be executed locally. .. GENERATED FROM PYTHON SOURCE LINES 178-187 .. code-block:: default if __name__ == "__main__": print("Running basic_boolean_wf a few times") for i in range(0, 5): print(f"Basic boolean wf output {basic_boolean_wf()}") print( f"Boolean input {True if i < 2 else False}, workflow output {bool_input_wf(b=True if i < 2 else False)}" ) .. GENERATED FROM PYTHON SOURCE LINES 188-193 Example 6 ^^^^^^^^^ It is possible to arbitrarily nest conditional sections inside other conditional sections. The conditional sections can only be in the ``then`` part of the previous conditional block. This example shows how float comparisons can be used to create a multi-level nested workflow. .. GENERATED FROM PYTHON SOURCE LINES 193-214 .. code-block:: default @workflow def nested_conditions(my_input: float) -> float: return ( conditional("fractions") .if_((my_input > 0.1) & (my_input < 1.0)) .then( conditional("inner_fractions") .if_(my_input < 0.5) .then(double(n=my_input)) .elif_((my_input > 0.5) & (my_input < 0.7)) .then(square(n=my_input)) .else_() .fail("Only <0.7 allowed") ) .elif_((my_input > 1.0) & (my_input < 10.0)) .then(square(n=my_input)) .else_() .then(double(n=my_input)) ) .. GENERATED FROM PYTHON SOURCE LINES 215-216 The nested conditionals can be executed locally. .. GENERATED FROM PYTHON SOURCE LINES 216-220 .. code-block:: default if __name__ == "__main__": print(f"nested_conditions(0.4) -> {nested_conditions(my_input=0.4)}") .. GENERATED FROM PYTHON SOURCE LINES 221-227 Example 7 ^^^^^^^^^^ It is possible to consume the outputs from conditional nodes. In the case of conditionals, the outputs are computed as a subset of outputs produced by ``then`` nodes. In this example, we call ``square()`` in one condition and ``double()`` in another. .. GENERATED FROM PYTHON SOURCE LINES 227-235 .. code-block:: default @task def calc_sum(a: float, b: float) -> float: """ returns the sum of a and b. """ return a + b .. GENERATED FROM PYTHON SOURCE LINES 236-246 Altogether, the workflow that consumes outputs from conditionals can be constructed as shown. .. tip:: A useful mental model to consume outputs of conditions is to think of them as ternary operators in programming languages. The only difference is that they can be n-ary. In Python, this is equivalent to .. code-block:: python x = 0 if m < 0 else 1 .. GENERATED FROM PYTHON SOURCE LINES 246-262 .. code-block:: default @workflow def consume_outputs(my_input: float, seed: int = 5) -> float: is_heads = coin_toss(seed=seed) res = ( conditional("double_or_square") .if_(is_heads.is_true()) .then(square(n=my_input)) .else_() .then(calc_sum(a=my_input, b=my_input)) ) # Regardless of the result, call ``double`` before # the variable `res` is returned. In this case, ``res`` stores the value of the ``square`` or ``double`` of the variable `my_input` return double(n=res) .. GENERATED FROM PYTHON SOURCE LINES 263-264 The workflow can be executed locally. .. GENERATED FROM PYTHON SOURCE LINES 264-271 .. code-block:: default if __name__ == "__main__": print( f"consume_outputs(0.4) with default seed=5. This should return output of calc_sum => {consume_outputs(my_input=0.4)}" ) print( f"consume_outputs(0.4, seed=7), this should return output of square => {consume_outputs(my_input=0.4, seed=7)}" ) .. rst-class:: sphx-glr-timing **Total running time of the script:** ( 0 minutes 0.000 seconds) .. _sphx_glr_download_auto_core_control_flow_conditions.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: conditions.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: conditions.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_