Simplify Dockerfile.

Remove dependencies on Beam, install tensorflow-gnn from pip.

PiperOrigin-RevId: 592560197

Dockerfile

@@ -15,51 +15,14 @@
 # Example of building the docker image locally:
 # docker build . -t tfgnn:latest
 #
-# You can then start an interactive shell with:
-# docker run -it --entrypoint /bin/bash tfgnn:latest
-# TODO:Change this to tensorflow/tensorflow:nightly-jupyter
-FROM ubuntu:focal
+# You can then start an interactive python interpreter shell and 
+# import tensorflow-gnn with:
+# docker run -it tfgnn:latest
+FROM python:3.9-slim
 # tzdata asks questions.
 ENV DEBIAN_FRONTEND="noninteractive"
 ENV TZ="America/New_York"
-RUN apt-get -y update \
-  && apt-get install -y apt-transport-https \
-  && apt-get install -y curl \
-  && apt-get install -y gnupg \
-  && curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > bazel.gpg \
-  && mv bazel.gpg /etc/apt/trusted.gpg.d/ \
-  && echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list \
-  && apt-get -y update \
-  && apt-get install -y --no-install-recommends \
-  build-essential \
-  pkg-config \
-  bazel \
-  python3.9 \
-  python3.9-dev \
-  python3.9-venv \
-  graphviz-dev \
-  graphviz
 
-COPY . /app
-
-# Set up venv to avoid root installing/running python
-ENV VIRTUAL_ENV=/opt/venv
-RUN python3.9 -m venv ${VIRTUAL_ENV}
-# Put the virtual environment on the path.
-ENV PATH="${VIRTUAL_ENV}/bin:/app/tensorflow_gnn:${PATH}"
 RUN pip3 install --upgrade pip
 
-# TODO(b/246860101) Remove pins when protobuf-4/tensorflow incompatibilities are resolved.
-# TF issue: https://github.com/tensorflow/tensorflow/issues/56077
-RUN pip3 install "google-api-core==1.31.0" "apache-beam[gcp]==2.38.0" httplib2 notebook ogb
-
-# Add `--no-cache-dir` if disk space is an issue.
-RUN python3 -m pip install /app
-
-# Install the apache beam sdk for local and dataflow runner support.
-COPY --from=apache/beam_python3.9_sdk /opt/apache/beam /opt/apache/beam
-
-# Set the default entry point to the apache beam binary for out-of-the box
-# compatability with dataflow. This entrypoint can be overriden to run other
-# tf-gnn programs within this environment.
-ENTRYPOINT [ "/opt/apache/beam/boot" ]
+RUN pip3 install "tensorflow-gnn==1.0.0" httplib2 notebook ogb