Tfx: Unable to import libcuda.so.1 when using TFX on GPUs

[sadeel]

I am running TFX on KFP. I added a GPU to my workload by doing the following.

def use_gpu():
  def _set_gpu_spec(task):
    task.set_gpu_limit(1)

pipeline_operator_funcs = kubeflow_dag_runner.get_default_pipeline_operator_funcs()
pipeline_operator_funcs.append(use_gpu())
  config = kubeflow_dag_runner.KubeflowDagRunnerConfig(
      pipeline_operator_funcs=pipeline_operator_funcs,
      kubeflow_metadata_config=kubeflow_dag_runner
      .get_default_kubeflow_metadata_config(),
      tfx_image=tfx_image,
  )
kubeflow_dag_runner.KubeflowDagRunner(config=config).run(pipeline)

However, when I run this on Kubeflow, I get the following errors:
2020-03-25 14:19:06.119947: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2020-03-25 14:19:06.119993: E tensorflow/stream_executor/cuda/cuda_driver.cc:351] failed call to cuInit: UNKNOWN ERROR (303)

This is using tfx==0.21.0, which uses tensorflow==2.1.0. Note, that if I run my workload on KFP without TFX (using KFP DSL), it runs on GPU.

Thanks

[sadeel]

The /usr/local/cuda directories seem to be missing from the published TFX image (comparing to a tensorflow docker container, e.g. tensorflow/tensorflow:2.2.0rc0-gpu-py3)

[numerology]

Hi @sadeel
I think your inspection is correct. This is because published TFX image only picks tensorflow but not tensorflow-gpu.

@zhitaoli For such use case, do we think it's desirable to also allow users to have something like tensorflow/tfx-gpu?

[valeriano-manassero]

I ran into same issue and I temporarily solved creating a custom image and declaring it in Kubeflow pipeline.

Maybe my repo can be interesting for an official TFX image implementation: https://github.com/valeriano-manassero/tfx-nvidia-gpu

[ialdencoots]

Hello, just curious if there's been any more thought on this front around publishing an image like tensorflow/tfx-gpu. It seems to me to be a fairly common use case to want to use tfx/kubeflow with GPUs.