https://prefect.io logo
Join Slack
Channels
Powered by
does anyone have experience setting up prefect+ten...
# ask-community
d
does anyone have experience setting up prefect+tensorflow + cuda gpus in GKE? We already run prefect in GKE using a custom docker image based off a prefect image, but i couldn't get tensorflow to work with cuda..so i'm now trying to build our image off a GCP ML image, adding in both the prefect stuff and our own custom dependencies. Not succeeded yet, so if anyone has any pointers do shout!
j
I am running in GKE and I don't use the prefect image at all for the jobs - just install prefect in the build - so maybe don't start from the prefect image? This is the docker file;
Copy code
FROM python:3.9.9

COPY . /opt/src

WORKDIR /opt/src
RUN pip install --upgrade pip && pip install poetry && poetry build && pip install ./dist/*.whl
I am doing some slight weirdness to get the code picked up - deployments path arg points to 
/usr/local/lib/python3.9/site-packages/
🤔 1
d
i'm hoping i can do a multistage build that starts like this:
Copy code
FROM prefecthq/prefect:2-python3.9 as prefect


FROM <http://gcr.io/deeplearning-platform-release/tf-gpu.2-6.py39|gcr.io/deeplearning-platform-release/tf-gpu.2-6.py39>
WORKDIR /opt/prefect
COPY --from=prefect /opt/prefect/ /opt/prefect/
but i don't think it's quite that simple
j
🤷 - I am not sure what is actually needed from the prefect image for the jobs - I use the image for my agent though...
I had a play around with this since we may want to do something similar - did you look into just starting a vertex AI job from a standard flow? I wonder if thats a simpler way of approaching the setup?
d
did you look into just starting a vertex AI job from a standard flow
no i didn't try that. I think I have now got it working in GKE as I wanted, using roughly your approach rather than copying from the prefect image. My docker file is roughly the following:
Copy code
FROM <http://gcr.io/deeplearning-platform-release/tf-gpu.2-6.py39|gcr.io/deeplearning-platform-release/tf-gpu.2-6.py39>

# this is basically how the official prefect image works if you look carefully  
ADD <https://github.com/krallin/tini/releases/download/v0.19.0/tini> /usr/bin/tini
RUN chmod +x /usr/bin/tini
CMD ["python", "-m", "prefect.engine"]
ENTRYPOINT ["/usr/bin/tini", "-g", "--"] 

RUN pip install {" ".join(f'"{dep}"' for dep in python_dependencies)}

COPY pipelines/flows_config.py /opt/prefect/flows/flows.py
COPY . /opt/prefect/my_project/

WORKDIR /opt/prefect/my_project/
ENV PYTHONPATH=/opt/my_project/
and i think it's working, but need to double check a few things.
❤️ 1
i think i might be able to use a base image rather than the tensorflow one specifically (see here).
update... this is what i have ended up with as my dockerfile
Copy code
FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

# setup python 3.9, using deadsnakes ppa - <https://launchpad.net/~deadsnakes>
ENV TZ=Europe/London
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
RUN apt update -y \
    && apt -y install software-properties-common wget gcc \
    && add-apt-repository -y ppa:deadsnakes/ppa \
    && apt install -y python3.9 python3.9-dev python3.9-distutils \
    && ln -s /usr/bin/python3.9 /usr/bin/python \
    && wget <https://bootstrap.pypa.io/get-pip.py> \
    && python get-pip.py

# Install the hefty ML dependencies
RUN pip install tensorflow==2.12.0            
RUN pip install torch --index-url <https://download.pytorch.org/whl/cu118>
RUN pip install sentence-transformers~=2.2.2

RUN pip install {" ".join(f'"{dep}"' for dep in python_dependencies)}

COPY pipelines/flows_config.py /opt/prefect/flows/flows.py
COPY . /opt/prefect/my_project/

WORKDIR /opt/prefect/my_project/
ENV PYTHONPATH=/opt/my_project/
it means i now have control over the versions of: • cuda/cudnn/ubuntu ...in so far as nvidia offers various combinations on docker hub • python • tensorflow (as long as it's compatible with the above) • pytroch (again, as long as it's compatible with the above, although i believe it doesn't need cudnn only cuda and not sure if cuda comes bundled with it or quite what). and the image doesn't contain too much cruft that i don't want. It was a very painful process to get here, but it does seem to do the trick
8 Views
Open in Slack
PreviousNext
Powered by