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

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

[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.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No