.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto/integrations/kubernetes/ray_example/ray_example.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        Click :ref:`here <sphx_glr_download_auto_integrations_kubernetes_ray_example_ray_example.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_integrations_kubernetes_ray_example_ray_example.py:


Ray Tasks
----------
Ray task allows you to run a Ray job on an existing Ray cluster or create a Ray cluster by using the Ray operator.


Let's get started with an example!

.. GENERATED FROM PYTHON SOURCE LINES 11-12

First, we load the libraries.

.. GENERATED FROM PYTHON SOURCE LINES 12-19

.. code-block:: default

    import typing

    import ray
    from flytekit import Resources, task, workflow
    from flytekitplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig



.. GENERATED FROM PYTHON SOURCE LINES 20-24

Ray Task
=========

We define a ray_example `remote function <https://docs.ray.io/en/latest/ray-core/tasks.html#tasks>`__ that will be executed asynchronously in the Ray cluster.

.. GENERATED FROM PYTHON SOURCE LINES 24-29

.. code-block:: default

    @ray.remote
    def f(x):
        return x * x



.. GENERATED FROM PYTHON SOURCE LINES 30-40

Defining a Ray Config
====================

We create a HeadNodeConfig and WorkerNodeConfig for the Ray job, and these config will be used by Ray operator to launch a Ray cluster before running the task.

* ``ray_start_params``: `RayStartParams <https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-start>`__ are the params of the start command: address, object-store-memory
* ``replicas``: Desired replicas of the worker group. Defaults to 1.
* ``group_name``: RayCluster can have multiple worker groups, and it distinguishes them by name
* ``runtime_env``: A `runtime environment <https://docs.ray.io/en/latest/ray-core/handling-dependencies.html#runtime-environments>`__ describes the dependencies your Ray application needs to run, and it's installed dynamically on the cluster at runtime.


.. GENERATED FROM PYTHON SOURCE LINES 40-47

.. code-block:: default

    ray_config = RayJobConfig(
        head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
        worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
        runtime_env={"pip": ["numpy", "pandas"]},  # or runtime_env="./requirements.txt"
    )



.. GENERATED FROM PYTHON SOURCE LINES 48-55

Defining a Ray Task
===================
We use `Ray job submission <https://docs.ray.io/en/latest/cluster/job-submission.html#job-submission-architecture>`__ to run our ray_example tasks.
ray_task will be called in the Ray head node, and f.remote(i) will be executed asynchronously on separate Ray workers

.. note::
   The Resources here is used to define the resource of worker nodes

.. GENERATED FROM PYTHON SOURCE LINES 55-61

.. code-block:: default

    @task(task_config=ray_config, limits=Resources(mem="2000Mi", cpu="1"))
    def ray_task(n: int) -> typing.List[int]:
        futures = [f.remote(i) for i in range(n)]
        return ray.get(futures)



.. GENERATED FROM PYTHON SOURCE LINES 62-66

Workflow
========

Finally we define a workflow to call the ``ray_workflow`` task.

.. GENERATED FROM PYTHON SOURCE LINES 66-71

.. code-block:: default

    @workflow
    def ray_workflow(n: int) -> typing.List[int]:
        return ray_task(n=n)



.. GENERATED FROM PYTHON SOURCE LINES 72-73

We can run the code locally wherein Flyte creates a standalone Ray cluster locally.

.. GENERATED FROM PYTHON SOURCE LINES 73-75

.. code-block:: default

    if __name__ == "__main__":
        print(ray_workflow(n=10))


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.000 seconds)


.. _sphx_glr_download_auto_integrations_kubernetes_ray_example_ray_example.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example


    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: ray_example.py <ray_example.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: ray_example.ipynb <ray_example.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_