Airflow with Virtual Environments

[Surya Iyer]

Hello Community,

Our team is looking to integrate TFX into our MLOPS platform. We have a distributed airflow cluster. 

I wanted to know if it is possible to run each component in an virtual/Conda environments using the Airflow Dag Runner. So. it may not be necessary to install TFX on every worker node as it may have some conflicts with other DAGs already deployed in the server.

Thank you so much for the product. TFX is perfect for our platform!

Best,

Surya.

[Robert Crowe]

Hi Surya,

Thanks!  For your virtual environment I'd recommend virtualenv rather than Conda, we've seen better results.  For containerization we've mainly been focusing on Kubernetes, and for Airflow we've only really created examples for native components, but I know that containerization in Airflow is a thing.  Hopefully it just works?

What company are you with?

thanks,

Robert

Robert Crowe | TensorFlow Developer Engineer | rober...@google.com  | @robert_crowe

--
You received this message because you are subscribed to the Google Groups "TensorFlow Extended (TFX)" group.
To unsubscribe from this group and stop receiving emails from it, send an email to tfx+uns...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/tfx/aae69f05-af3d-4075-9a9e-3c4d9ee2e039n%40tensorflow.org.

[Surya Iyer]

Hi Robert,

I am working at Medallia.

Thanks for the information!  I do know that airflow as a whole can work in a virtual environment. But I am looking to run each component or the tfx DAG in a virtual environment . The process would look like ->

1. Take the requirements.txt

2. Create the environment

3. Run the component

4. Destroy the environment

We cannot install tfx as is in Airflow as there would be dependency conflicts with other processes which are in place.

Correct me if I am wrong.  Looking at the code base, I don't see a way to pass on the parameters for the virtual environment in the Airflow Dag Runner. And the airflow components  are child classes of Python Operator. Would I need to change this such that it inherits from Airflow's  PythonVirtualenvOperator?

To get the components to work in the virtual environment, is there a way to do it without changing the tfx source code?

Or would you suggest building a separate airflow cluster for tfx (may not be feasible)

Do you have any suggestions regarding best practices for tfx with airflow?  

Thanks a lot for your guidance!

--

Best,

Surya Iyer.

[Robert Crowe]

+Zhitao Li 

I think Zhitao may have some ideas.

Robert Crowe | TensorFlow Developer Engineer | rober...@google.com  | @robert_crowe

[Zhitao Li]

Hi Surya,

One idea is to run a cluster of Airflow workers in docker, in which you can easily have different dependencies. I personally found it more reliable than virtualenv for a production deployment.

I searched online and found https://airflow.apache.org/docs/apache-airflow/stable/start/docker.html. Maybe this would be a direction we could try?

--

Cheers,

Zhitao Li

[Gerard Casas Saez]

There are several ways to achieve isolation on Airflow. Probs the best recommendation for production is setting up the KubernetesExecutor and then using the pod_overwrite to change what image you want to use to execute each component of your DAG (via default_args in TFX AirflowConfig) and then point it to a docker image w the right environment. This is probs the best mode: https://airflow.apache.org/docs/apache-airflow/stable/executor/kubernetes.html#pod-override

If your deployment is Celery based, then you can achieve something similar by creating a separate queue that is set up with the right environment and making sure the queue has the right environment installed - https://airflow.apache.org/docs/apache-airflow/stable/executor/celery.html#queues

That said, Airflow community is really welcoming and would encourage you to join and ask if needed there (Slack and user mailing lists are usually great start for things like this). TFX specifically really just uses a derived PythonOpertor from Airflow to run its components, so this can be seen as more generally asking how to have isolated environments in Airflow, which there are several options available :D

Best,

Gerard Casas Saez
Twitter | Cortex | @casassaez

To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/tfx/CALRp7OrAvrPrpdktJD%2BgH52ZrcPodrs9me8%3DF_-baF2k%3DtchJA%40mail.gmail.com.