diff --git a/inference/maxengine_server/Dockerfile b/inference/maxengine_server/Dockerfile new file mode 100644 index 000000000..b94ea499a --- /dev/null +++ b/inference/maxengine_server/Dockerfile @@ -0,0 +1,46 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Ubuntu:22.04 +# Use Ubuntu 22.04 from Docker Hub. +# https://hub.docker.com/_/ubuntu/tags?page=1&name=22.04 +FROM ubuntu:22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV MAXTEXT_VERSION=main + +RUN apt -y update && apt install -y --no-install-recommends \ + ca-certificates \ + git \ + python3.10 \ + python3-pip + +RUN update-alternatives --install \ + /usr/bin/python3 python3 /usr/bin/python3.10 1 + +RUN git clone https://github.com/AI-Hypercomputer/maxtext.git && \ +git clone https://github.com/AI-Hypercomputer/JetStream.git + +RUN cd maxtext/ && \ +git checkout ${MAXTEXT_VERSION} && \ +bash setup.sh + +RUN cd /JetStream && \ +pip install -e . + +COPY maxengine_server_entrypoint.sh /usr/bin/ + +RUN chmod +x /usr/bin/maxengine_server_entrypoint.sh + +ENTRYPOINT ["/usr/bin/maxengine_server_entrypoint.sh"] \ No newline at end of file diff --git a/inference/maxengine_server/README.md b/inference/maxengine_server/README.md new file mode 100644 index 000000000..40fad64b5 --- /dev/null +++ b/inference/maxengine_server/README.md @@ -0,0 +1,14 @@ +## Build and upload Maxengine Server image + +These instructions are to build the Maxengine Server image, which calls an entrypoint script that invokes the [JetStream](https://github.com/AI-Hypercomputer/JetStream) inference server with the MaxText framework. + +``` +docker build -t maxengine-server . +docker tag maxengine-server us-docker.pkg.dev/${PROJECT_ID}/jetstream/maxengine-server:latest +docker push us-docker.pkg.dev/${PROJECT_ID}/jetstream/maxengine-server:latest +``` + +If you would like to change the version of MaxText the image is built off of, change the `MAXTEXT_VERSION` environment variable: +``` +ENV MAXTEXT_VERSION= +``` \ No newline at end of file diff --git a/inference/maxengine_server/maxengine_server_entrypoint.sh b/inference/maxengine_server/maxengine_server_entrypoint.sh new file mode 100644 index 000000000..726926ce0 --- /dev/null +++ b/inference/maxengine_server/maxengine_server_entrypoint.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +cd /maxtext +python3 MaxText/maxengine_server.py \ +MaxText/configs/base.yml $@ \ No newline at end of file