Using custom code on a Coiled cluster
=====================================

In this tutorial, you'll learn more advanced techniques for installing custom Python code on an already running cluster, using a custom Docker image to create a Coiled software environment, or running additional commands after installing a package.

Replicate your local Python environment
---------------------------------------

Package sync scans your local Python environment and replicates it on the cluster—even local packages and Git dependencies. It’s easier than building a Docker image, plus it launches significantly faster (see :doc:`../package_sync`). Package sync is used by default when you create a cluster.

Install a local Python module
-----------------------------

If you're not using package sync, but you have custom Python packages you would like to use on your Coiled cluster, you can upload a Python module to all of your workers using Dask's `upload_file <https://distributed.dask.org/en/latest/api.html?highlight=upload_file#distributed.Client.upload_file>`_. This sends a local file to all your worker nodes on an already running cluster. This method is particularly helpful for when you would like to use custom Python functions that haven't been packaged in your Cluster.

.. code:: python

    import coiled
    from dask.distributed import Client

    cluster = coiled.Cluster()
    client = Client(cluster)

    client.upload_file("my_script.py")


.. _custom-docker:

.. include:: ../docker.rst

Install pip-installable packages
--------------------------------

If you have a package that is pip-installable, but not yet publicly available on PyPI or conda-forge, for example, you can use Dask's `PipInstall Worker Plugin <https://distributed.dask.org/en/latest/plugins.html?highlight=PipInstall#built-in-worker-plugins>`_ to pip install a set of packages. This is particularly useful for uploading modules that are still in development.

You can upload a public module in GitHub

.. code-block:: python

  from dask.distributed.diagnostics.plugin import PipInstall

  plugin = PipInstall(packages=["git+<github url>"])
  client.register_worker_plugin(plugin, name="<dependency name>")

If you want to install from a private repository you need to have a GitHub token set in your account by either having signed up with GitHub or by :doc:`adding your GitHub token to your profile <github_tokens>`. GitHub tokens are stored with Coiled and then used on the machine that's building the software environment; the token is not saved in the software environment.

.. code-block:: python

  from dask.distributed.diagnostics.plugin import PipInstall

  plugin = PipInstall(packages=["git+https://{GITHUB_TOKEN}@github.com/<repo>"])
  client.register_worker_plugin(plugin, name="<dependency name>")

.. note::

   Using the ``name=`` argument will allow you to call ``PipInstall`` more than once, otherwise you might see a message from workers like ``{'tls://10.4.1.170:38403': {'status': 'repeat'}``.

Upload a local directory
------------------------

Similar to the ``PipInstall`` Plugin, you can upload a local directory to your cluster by using the `UploadDirectory Nanny Plugin <https://distributed.dask.org/en/latest/plugins.html?highlight=PipInstall#built-in-nanny-plugins>`_.

You can upload a local directory from your machine to the cluster using:

.. code-block:: python

  from distributed.diagnostics.plugin import UploadDirectory

  client.register_worker_plugin(UploadDirectory("/path/to/directory"), nanny=True)

It's worth noting ``UploadDirectory`` will not install anything on its own. Ideally, you would package the code and directly use the ``PipInstall`` Worker Plugin mentioned above. However, if this is not possible you can write your own plugin for uploading and installation:

.. code-block::

    class InstallModulePlugin(UploadDirectory):
    """Use this plugging to upload a directory and install that directory in the workers."""
    def __init__(self, dir_path, module):
        """Initializes the plugin

        Arguments
        ---------
        dir_path: str, path to the directory you want to upload
        module: directory name
        """
        super().__init__(dir_path, update_path=True)
        self.module = module

    async def setup(self, nanny):
        await super().setup(nanny)

        import os
        import subprocess
        path_to_module = os.path.join(nanny.local_directory, self.module)

        # or whatever bash command to install package
        subprocess.call(["pip", "install", "-e", path_to_module])


    plugin = InstallModulePlugin("path_to_directory", "directory_name")
    client.register_worker_plugin(plugin, nanny=True)
