FROM pytorch/pytorch:2.7.1-cuda12.6-cudnn9-devel

# torch + triton are already in this image - do not reinstall torch via pip.
# Verify triton: python -c "import triton; print(triton.__version__)"

WORKDIR /workspace

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    curl \
    ffmpeg \
    git \
    libasound2 \
    libdbus-glib-1-2 \
    libdrm-dev \
    libgbm-dev \
    libnss3 \
    libpulse-dev \
    libx11-6 \
    libxcb1 \
    libxcomposite1 \
    libxcursor1 \
    libxdamage1 \
    libxext6 \
    libxfixes3 \
    libxi6 \
    libxkbcommon-dev \
    libxrandr2 \
    libxrender1 \
    libxshmfence-dev \
    ninja-build \
    openjdk-17-jdk-headless \
    sudo \
    unzip \
    wget \
    xauth \
    xvfb \
    && rm -rf /var/lib/apt/lists/*

# Android SDK / emulator for AndroidWorld live RL.
ARG ANDROID_API_LEVEL=33
ARG ANDROID_BUILD_TOOLS=33.0.0
ARG ANDROID_CMDLINE_TOOLS=commandlinetools-linux-11076708_latest.zip
ARG ANDROID_SDK_ROOT=/opt/android
ARG ANDROID_EMULATOR_NAME=Pixel_6_API_33
ARG ANDROID_EMULATOR_DEVICE=pixel_6
ARG ANDROID_SYSTEM_IMAGE=system-images;android-${ANDROID_API_LEVEL};google_apis;x86_64
ENV ANDROID_SDK_ROOT=${ANDROID_SDK_ROOT}
ENV ANDROID_HOME=${ANDROID_SDK_ROOT}
ENV ANDROID_AVD_HOME=/root/.android/avd
ENV EMULATOR_NAME=${ANDROID_EMULATOR_NAME}
ENV DEVICE_NAME=${ANDROID_EMULATOR_DEVICE}
ENV PATH="${PATH}:${ANDROID_SDK_ROOT}/cmdline-tools/latest/bin:${ANDROID_SDK_ROOT}/emulator:${ANDROID_SDK_ROOT}/platform-tools:${ANDROID_SDK_ROOT}/build-tools/${ANDROID_BUILD_TOOLS}"

RUN mkdir -p ${ANDROID_SDK_ROOT}/cmdline-tools /root/.android && \
    wget -q https://dl.google.com/android/repository/${ANDROID_CMDLINE_TOOLS} -O /tmp/android_cmdline_tools.zip && \
    unzip -q /tmp/android_cmdline_tools.zip -d /tmp/android_cmdline_tools && \
    mv /tmp/android_cmdline_tools/cmdline-tools ${ANDROID_SDK_ROOT}/cmdline-tools/latest && \
    rm -rf /tmp/android_cmdline_tools /tmp/android_cmdline_tools.zip && \
    yes | sdkmanager --licenses >/dev/null && \
    sdkmanager --no_https \
      "platform-tools" \
      "emulator" \
      "platforms;android-${ANDROID_API_LEVEL}" \
      "build-tools;${ANDROID_BUILD_TOOLS}" \
      "${ANDROID_SYSTEM_IMAGE}" && \
    echo "no" | avdmanager create avd --force \
      --name "${ANDROID_EMULATOR_NAME}" \
      --device "${ANDROID_EMULATOR_DEVICE}" \
      --package "${ANDROID_SYSTEM_IMAGE}"

# Install all deps except torch, which is already in the base image.
RUN pip install --no-cache-dir \
    "transformers>=4.50.0" \
    "datasets>=2.20.0" \
    "Pillow>=10.0.0" \
    "bitsandbytes>=0.43.0" \
    "peft>=0.11.0" \
    "accelerate>=0.30.0" \
    "pyarrow>=16.0.0" \
    "tqdm>=4.66.0" \
    "numpy>=1.26.0" \
    "sentencepiece>=0.2.0" \
    "packaging>=24.0" \
    "setuptools>=69.0" \
    "wheel>=0.43.0"

# Optional fast attention path for Qwen/Gated Delta Net style models.
# Build this in Linux so Windows CUDA toolkit state does not matter.
ENV MAX_JOBS=4
ENV TORCH_CUDA_ARCH_LIST="8.6"
RUN pip install --no-cache-dir --no-deps "flash-linear-attention==0.5.0" "fla-core==0.5.0" \
    && pip install --no-cache-dir --no-build-isolation "causal-conv1d"

RUN pip install --no-cache-dir "einops>=0.8.0"

# AndroidWorld is installed into the image so a collaborator can run the
# emulator/server without host Python setup or a checked-out external/ tree.
ARG ANDROID_WORLD_REPO=https://github.com/google-research/android_world.git
ARG ANDROID_WORLD_COMMIT=d9c569f764b3a5629321858de03ff653d0f24056
RUN git clone ${ANDROID_WORLD_REPO} /opt/android_world && \
    cd /opt/android_world && \
    git checkout ${ANDROID_WORLD_COMMIT}
RUN chmod +x /opt/android_world/docker_setup/start_emu.sh \
    /opt/android_world/docker_setup/start_emu_headless.sh \
    /opt/android_world/docker_setup/entrypoint.sh && \
    pip install --no-cache-dir --no-build-isolation -e /opt/android_world && \
    cd /opt/android_world && \
    python setup.py generate_protos

# The upstream Docker server performs AndroidWorld app setup on every startup.
# Keep it available, but make it opt-in so the server can boot as a reusable
# emulator/backend for our own harness.
RUN python - <<'PY'
from pathlib import Path
p = Path("/opt/android_world/server/android_server.py")
s = p.read_text()
if "import os\n" not in s:
    s = s.replace("import contextlib\n", "import contextlib\nimport os\n")
s = s.replace(
    "      emulator_setup=True,\n",
    "      emulator_setup=os.environ.get('ANDROID_WORLD_EMULATOR_SETUP', '0') == '1',\n",
)
p.write_text(s)
PY

# Contacts setup can fail on Android 13 images when the notification prompt is
# absent or has already been resolved. Treat that prompt as optional; the task
# state is still usable as long as Contacts launches.
RUN python - <<'PY'
from pathlib import Path
p = Path("/opt/android_world/android_world/env/setup_device/apps.py")
s = p.read_text()
old = '''      # Allow Contacts to send you notifications?
      controller.click_element("Don't allow")
      time.sleep(2.0)
'''
new = '''      # Allow Contacts to send you notifications? This dialog is image/build
      # dependent and may already be resolved on Android 13 emulator images.
      try:
        controller.click_element("Don't allow")
        time.sleep(2.0)
      except ValueError:
        logging.info("Contacts notification permission prompt not present.")
'''
if old not in s:
    raise SystemExit("Contacts setup patch target not found")
p.write_text(s.replace(old, new))
PY

# Source is mounted as a volume at runtime - no COPY needed.
# This means code changes do not require a rebuild.

# HF_TOKEN is injected at runtime via --env or docker-compose.
ENV PYTHONUNBUFFERED=1

CMD ["python", "-u", "scripts/train_a11y.py"]