Airflow and Rapids integration in docker not working

Accelerated Computing CUDA CUDA Setup and Installation
, , ,
1

Hi: I am building a docker image with rapids and airflow but I am having issues with the integration. When running a dag in airflow that contains cudf code the system throws and error message.

Here is the error message:

Traceback (most recent call last):
  File "/opt/airflow/venv/lib/python3.11/site-packages/airflow/models/taskinstance.py", line 767, in _execute_task
    result = _execute_callable(context=context, **execute_callable_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/airflow/models/taskinstance.py", line 733, in _execute_callable
    return ExecutionCallableRunner(
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/airflow/utils/operator_helpers.py", line 252, in run
    return self.func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/airflow/models/baseoperator.py", line 406, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/airflow/operators/python.py", line 238, in execute
    return_value = self.execute_callable()
                   ^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/airflow/operators/python.py", line 256, in execute_callable
    return runner.run(*self.op_args, **self.op_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/airflow/utils/operator_helpers.py", line 252, in run
    return self.func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/dags/PIS/PIS_prepare_training_file.py", line 20, in get_training_file
    s = cudf.Series([1, 2, 3, None, 4, 5])
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/cudf/utils/performance_tracking.py", line 51, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/cudf/core/series.py", line 672, in __init__
    column = as_column(
             ^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/cudf/core/column/column.py", line 2181, in as_column
    return as_column(arbitrary, nan_as_null=nan_as_null, dtype=dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/cudf/core/column/column.py", line 1808, in as_column
    col = ColumnBase.from_arrow(arbitrary)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/airflow/venv/lib/python3.11/site-packages/cudf/core/column/column.py", line 368, in from_arrow
    result = libcudf.interop.from_arrow(data)[0]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "interop.pyx", line 163, in cudf._lib.interop.from_arrow
  File "/usr/lib/python3.11/functools.py", line 909, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "interop.pyx", line 134, in pylibcudf.interop._from_arrow_table
RuntimeError: Fatal CUDA error encountered at: /__w/cudf/cudf/cpp/src/bitmask/null_mask.cu:93: 3 cudaErrorInitializationError initialization error

However when running the container in interactive mode I can see the output.

docker exec -it 747543630bc09099812ffa60e442dac9bb4245fdd4f178e5eaa88389c073b664 python -c "import cudf; print(cudf.Series([1, 2, 3]))"

$ ls
NGC-DL-CONTAINER-LICENSE  bin  boot  dev  entrypoint  etc  home  lib  lib.usr-is-merged  lib64  media  mnt  opt  proc  requirements.txt  root  run  sbin  srv  sys  tmp  usr  var
$ python
Python 3.11.10 (main, Sep  7 2024, 18:35:41) [GCC 13.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import cudf
>>> s = cudf.Series([1, 2, 3, None, 4, 5])
>>> s
0       1
1       2
2       3
3    <NA>
4       4
5       5
dtype: int64

My airflow version is: 2.10.2
Python is: 3.11
Rapids is: 24.10
Cuda: 12.5

My docker file starts with nvidia/cuda:12.5.1-base-ubuntu24.04

image
image1732×782 21.2 KB

ENV:

ENV PATH=/usr/local/cuda-${CUDA_VERSION}/bin${PATH:+:${PATH}}

ENV LD_LIBRARY_PATH=/usr/local/cuda-${CUDA_VERSION}/lib64

Any ideas of how to make it work in an airflow dag?

Thanks in advance