Update blake3 dep, remove dockerfile, simplify build (#1561)

Author: nnarayen

context_builder.Dockerfile

@@ -5,24 +5,14 @@
 FROM python:3.9-slim AS builder
 
 RUN apt-get update \
-    && apt-get install --yes --no-install-recommends curl build-essential golang-go \
+    && apt-get install --yes --no-install-recommends curl \
     && apt-get autoremove -y \
     && apt-get clean -y \
     && rm -rf /var/lib/apt/lists/* /tmp/library-scripts/
 
 RUN curl -sSL https://install.python-poetry.org | python -
 ENV PATH="/root/.local/bin:${PATH}"
 
-# NB(nikhil): We need to explicitly install Rust and golang-go on arm64 architectures
-# for certain truss dependencies (blake3, dockerfile) that don't have prebuilt wheels.
-# This is a multi stage image, so the extra tools don't bloat the final image.
-RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
-ENV PATH="/root/.cargo/bin:${PATH}"
-RUN rustup update stable && rustup default stable
-
-# Will only exist for arm64 builds, but is a no-op otherwise.
-ENV PATH="/root/.cargo/bin:${PATH}"
-
 # Chains/train source code is not actually needed (and deps from chains group won't be
 # installed when using `--only builder`). But nonetheless poetry fails the install
 # if this directory is not present/empty - so we copy it.

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.78"
+version = "0.9.79rc001"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"
@@ -85,7 +85,7 @@ single-source = "^0.3.0"
 Jinja2 = { version = "^3.1.2", optional = false }
 aiofiles = { version = "^24.1.0", optional = false }
 aiohttp = { version = "^3.10.10", optional = false }
-blake3 = { version = "^0.3.3", optional = false }
+blake3 = { version = "^1.0.4", optional = false }
 boto3 = { version = "^1.34.85", optional = false }
 click = { version = "^8.0.3", optional = false }
 fastapi = { version =">=0.109.1", optional = false }
@@ -159,7 +159,6 @@ types-setuptools = "^69.0.0.0"
 [tool.poetry.group.dev-server.dependencies]
 # These packages are needed to run local tests of server components. Note that the actual
 # server deps for building the docker image are (so far) defined in `requirements.txt`-files.
-dockerfile = "^3.2.0"
 fastapi =">=0.109.1"
 flask = "^2.3.3"
 msgpack = ">=1.0.2"

pyproject.toml

@@ -10,8 +10,6 @@
 HUGGINGFACE_TRANSFORMER = "huggingface_transformer"
 LIGHTGBM = "lightgbm"
 
-ARM_PLATFORMS = ("aarch64", "arm64")
-
 
 _TRUSS_ROOT = pathlib.Path(__file__).parent.parent.resolve()
 

truss/base/constants.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-import platform
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
@@ -17,7 +16,6 @@
 
 from truss.base import constants
 from truss.base.constants import (
-    ARM_PLATFORMS,
     BASE_SERVER_REQUIREMENTS_TXT_FILENAME,
     BEI_MAX_CONCURRENCY_TARGET_REQUESTS,
     BEI_REQUIRED_MAX_NUM_TOKENS,
@@ -711,8 +709,6 @@ def _render_dockerfile(
             MIN_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE.split(".")[0]
         )
 
-        should_install_arm64_requirements = platform.machine() in ARM_PLATFORMS
-
         hf_access_token = config.secrets.get(constants.HF_ACCESS_TOKEN_KEY)
         dockerfile_contents = dockerfile_template.render(
             should_install_server_requirements=should_install_server_requirements,
@@ -745,7 +741,6 @@ def _render_dockerfile(
             external_data_files=external_data_files,
             build_commands=build_commands,
             use_local_src=config.use_local_src,
-            should_install_arm64_requirements=should_install_arm64_requirements,
             **FILENAME_CONSTANTS_MAP,
         )
 

truss/contexts/image_builder/serving_image_builder.py

@@ -2,6 +2,10 @@
 from pathlib import Path
 from typing import Dict, List
 
+from truss.contexts.local_loader.dockerfile_parser import (
+    DockerInstruction,
+    parse_dockerfile,
+)
 from truss.util.path import copy_tree_or_file
 
 
@@ -23,9 +27,7 @@ class DockerBuildEmulator:
     """
 
     def __init__(self, dockerfile_path: Path, context_dir: Path) -> None:
-        import dockerfile
-
-        self._commands = dockerfile.parse_file(str(dockerfile_path))
+        self._commands = parse_dockerfile(dockerfile_path)
         self._context_dir = context_dir
 
     def run(self, fs_root_dir: Path) -> DockerBuildEmulatorResult:
@@ -41,18 +43,20 @@ def _resolve_values(keys: List[str]) -> List[str]:
 
         result = DockerBuildEmulatorResult()
         for cmd in self._commands:
-            if cmd.cmd not in ["ENV", "ENTRYPOINT", "COPY", "WORKDIR"]:
+            if not cmd.is_supported:
                 continue
             values = _resolve_values(cmd.value)
-            if cmd.cmd == "ENV":
+            if cmd.instruction == DockerInstruction.ENV:
+                if "=" in values[0]:
+                    values = values[0].split("=", 1)
                 result.env[values[0]] = values[1]
-            if cmd.cmd == "ENTRYPOINT":
+            if cmd.instruction == DockerInstruction.ENTRYPOINT:
                 result.entrypoint = list(values)
-            if cmd.cmd == "COPY":
+            if cmd.instruction == DockerInstruction.COPY:
                 src, dst = values
                 src = src.replace("./", "", 1)
                 dst = dst.replace("/", "", 1)
                 copy_tree_or_file(self._context_dir / src, fs_root_dir / dst)
-            if cmd.cmd == "WORKDIR":
+            if cmd.instruction == DockerInstruction.WORKDIR:
                 result.workdir = result.workdir / values[0]
         return result

truss/contexts/local_loader/docker_build_emulator.py

@@ -0,0 +1,52 @@
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import List
+
+
+class DockerInstruction(Enum):
+    COPY = "COPY"
+    ENV = "ENV"
+    ENTRYPOINT = "ENTRYPOINT"
+    WORKDIR = "WORKDIR"
+    UNKNOWN = "UNKNOWN"
+
+    @classmethod
+    def from_string(cls, instruction: str) -> "DockerInstruction":
+        try:
+            return cls(instruction.upper())
+        except ValueError:
+            return cls.UNKNOWN
+
+
+@dataclass
+class DockerCommand:
+    instruction: DockerInstruction
+    value: List[str]
+
+    @property
+    def is_supported(self) -> bool:
+        return self.instruction != DockerInstruction.UNKNOWN
+
+
+def parse_dockerfile(dockerfile_path: Path) -> List[DockerCommand]:
+    commands = []
+    with open(dockerfile_path, "r") as f:
+        content = f.read()
+
+    # Join across line continuations.
+    content = content.replace("\\\n", " ")
+    for line in content.splitlines():
+        line = line.strip()
+
+        if not line or line.startswith("#"):
+            continue
+
+        parts = line.split(" ", 1)
+        if len(parts) == 2:
+            instruction_str, args_str = parts
+            instruction = DockerInstruction.from_string(instruction_str)
+            args = args_str.split()
+            commands.append(DockerCommand(instruction=instruction, value=args))
+
+    return commands

truss/contexts/local_loader/dockerfile_parser.py

@@ -25,20 +25,6 @@ RUN pip install mkl
 {% block post_base %}
 {% endblock %}
 
-# NB(nikhil): We need to install additional toolchains to support installing
-# python dependencies (currently blake3 and dockerfile) on arm64 architecture.
-# Currently, we only support arm64 in local dev, but we need these for:
-#   1) Control server (depends on truss)
-#   2) Python DX / Chains (depends on truss)
-{% block install_arm64_requirements %}
-    {%- if should_install_arm64_requirements %}
-RUN apt-get update && apt-get install --yes --no-install-recommends golang-go
-RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
-ENV PATH="/root/.cargo/bin:${PATH}"
-RUN rustup update stable && rustup default stable
-    {%- endif %}
-{% endblock %}
-
 {% block install_system_requirements %}
     {%- if should_install_system_requirements %}
 COPY ./{{system_packages_filename}} {{system_packages_filename}}

truss/templates/base.Dockerfile.jinja

@@ -1,3 +1,2 @@
 torch>=2.0.1
 blake3>=0.3.3
-dockerfile>=3.2.0