I'm having some trouble getting GPUs working on AWS kubernetes. I've got the nvidia-device-plugin installed and my jobs are scheduling on my GPU instance when using

@kubernetes(gpu=1)

I'm getting stuck trying to install

tensorflow-gpu

. When running the following code:

from metaflow import FlowSpec, step, kubernetes, retry, conda_base

libraries = {
    "tensorflow-gpu": "2.11.1",
    "cudatoolkit": "11.0.3"
}

@conda_base(libraries=libraries, python="3.10.9")
class Tensorflow(FlowSpec):

Step: start, Error: command '['/opt/conda/condabin/mamba', 'create', '--yes', '--no-default-packages', '--name', 'metaflow_Tensorflow_linux-64_615c884d84bb1f7ae93a2e331b5189159b559cc8', '--quiet', b'python==3.10.9', b'requests==>=2.21.0', b'boto3==>=1.14.0', b'tensorflow-gpu==2.11.1', b'cudatoolkit==11.0.3']' returned error (1): b'Could not solve for environment specs\nEncountered problems while solving:\n  - nothing provides __cuda needed by tensorflow-2.11.1-cuda112py310he87a039_0\n\nThe environment can\'t be solved, aborting the operation\n\n{\n    "success": false\n}\n', stderr=b''

I was able to shell into the running docker container, install tensorflow-gpu manually, and it worked fine. Any thoughts?

Use envvar CONDA_OVERRIDE_CUDA=11.4 (or whatever).

I tried that to no avail:

$ CONDA_CUDA_OVERRIDE="11.8" METAFLOW_DEFAULT_CONTAINER_IMAGE=<http://075106392831.dkr.ecr.us-east-1.amazonaws.com/test:tensorflow|075106392831.dkr.ecr.us-east-1.amazonaws.com/test:tensorflow> python k8s_basic.py --environment=conda run
Metaflow 2.8.3 executing ForeachFlow for user:tyler.potts
Validating your flow...
    The graph looks good!
Running pylint...
    Pylint is happy!
Bootstrapping conda environment...(this could take a few minutes)
    Conda ran into an error while setting up environment.:
    Step: start, Error: command '['/opt/conda/condabin/mamba', 'create', '--yes', '--no-default-packages', '--name', 'metaflow_ForeachFlow_linux-64_6e0539c8f0fc87f5cf2003216d0053a447c1be45', '--quiet', b'python==3.10.9', b'requests==>=2.21.0', b'boto3==>=1.14.0', b'tensorflow-gpu==2.11.1', b'cudatoolkit==11.8.0', b'cudnn==8.8.0.121']' returned error (1): b'Could not solve for environment specs\nEncountered problems while solving:\n  - nothing provides __cuda needed by tensorflow-2.11.1-cuda112py310he87a039_0\n\nThe environment can\'t be solved, aborting the operation\n\n{\n    "success": false\n}\n', stderr=b''

I even went so far as to add the variable into the conda.py file and pip install a local version: https://github.com/Netflix/metaflow/blob/master/metaflow/plugins/conda/conda.py#L179

"MAMBA_JSON": "True",
                "CONDA_CUDA_OVERRIDE": "11.2"

try adding

@environment(vars={'CONDA_CUDA_OVERRIDE': '11.2'})

so the env var gets passed to the container

note it’s

CONDA_OVERRIDE_CUDA=11.4

not

CONDA_CUDA_OVERRIDE=11.4

. See: https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-virtual.html

@dry-beach-38304 @straight-shampoo-11124 Thanks for the direction I indeed had my cuda override flag incorrect. It worked with that flag, and this is the dockerfile I built that worked with gpus for future posterity (liberally copied from other comments around here):

FROM tensorflow/tensorflow:2.12.0-gpu

RUN apt update && apt install -y curl python3 python3-pip wget

RUN wget <https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh> -O ~/miniconda.sh
RUN bash ~/miniconda.sh -b
RUN rm ~/miniconda.sh

RUN /root/miniconda3/bin/conda init && cp /root/.bashrc /root/.bash_profile

RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 10 && \
    update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 10

ENV LC_ALL=C.UTF-8
ENV LANG=C.UTF-8

oh nice! Good to hear that you got it working 🤗