nandi/frq · _loader.py

nandi/frqpublic⑂ Fork 0
⑂ 1403956859648c3d39892e3994df8fee863ec653
_loader.py · 476 lines · 21.3 KBPython    Blame HistoryRaw
Build the window somewhere with room for it   1"""Turn a `container.toml` into a `modal.Image` and a `modal.App`.
Carry the loader's .git fix over   3Every container -- under `containers/` here, `.modal/` in a repo that merely
4builds itself on Modal -- is a directory with a `container.toml` and
Build the window somewhere with room for it   5a stub `container.py`. See `spec.md` for the keys; this file is what reads
6them. Nothing here is Modal-specific configuration in its own right -- each
7spec key maps onto a documented Modal argument, and the mapping is meant to
8stay boring enough to read straight through.
11
12import shlex
13import subprocess
14import tomllib
16import modal
18REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
19PTYSHIM_C = os.path.join(REPO, "ptyshim.c")
20SHIM_SO = "/opt/ptyshim.so"
23class SpecError(Exception):
  """The container.toml says something that cannot be built."""
27def _running_in_modal() -> bool:
  """True inside a Modal container, false on the machine that launched it."""
  return bool(os.environ.get("MODAL_TASK_ID"))
32class Container:
  """One container: its spec, its image, its app, and how to run it."""
  def __init__(self, spec: dict, directory: str):
      self.is_remote = _running_in_modal()
      self.dir = directory
      self.spec = spec
      self.name = self._require("container", "name")
      self.description = spec.get("container", {}).get("description", "")
      run = spec.get("run", {})
      self.workdir = run.get("workdir", "/app")
      self.command = run.get("command", "")
      self.env = dict(run.get("env", {}))
A shell that is already the devShell, and two recipes to reach it   47        # Substituters reach nix through NIX_CONFIG rather than nix.conf, and
      # at run time rather than build time. Written into the image's nix.conf
      # instead, nix touches the path while the image is being built and
      # creates it -- and Modal will not mount a volume over a non-empty
      # directory, so the container that wanted the cache cannot start. As
      # env it covers every nix command that runs in the container, the ones
      # typed by hand in a shell included, and leaves the mount point empty.
      if subs := list(spec.get("nix", {}).get("substituters", [])):
          line = "extra-substituters = " + " ".join(subs)
          self.env["NIX_CONFIG"] = (
              f"{self.env['NIX_CONFIG']}\n{line}" if "NIX_CONFIG" in self.env
              else line
Build the window somewhere with room for it   61        # "function" (the default) or "sandbox". Sandboxes can run on a real
      # VM, which Functions cannot -- see ../README.md.
      self.runtime = spec.get("container", {}).get("runtime", "function")
      nix = spec.get("nix", {})
      self.use_flake = bool(nix.get("flake", False))
      self.use_shim = bool(nix.get("shim", False))
      # name -> mount path. Modal Volumes, mounted while the container runs
      # and NOT while its image is built: a volume mount is not part of the
      # resulting image, so anything written to one during a build step is
      # gone by the time the container starts. Persist across runs is the
      # whole point -- a nix store to substitute from, a cargo target
      # directory, a dataset too big to bake in.
      self.volume_spec = dict(spec.get("volumes", {}))
A third target, and the seam that was already waiting for it   77        # Ports to tunnel out of a Sandbox, from [network] ports. Encrypted,
      # which is Modal's own word for "the tunnel terminates TLS and speaks
      # plain HTTP to your process" -- so the thing listening inside is an
      # ordinary http.server and not something holding a certificate.
      # Sandbox-only: a Function has no long-lived process to tunnel into,
      # and `modal run` on one would hand back a URL for a container that
      # has already exited.
      self.ports = [int(p) for p in spec.get("network", {}).get("ports", [])]
Build the window somewhere with room for it   87        # Modal re-imports this module inside the container, so everything
      # below runs twice: once here, once out there. Out there the local
      # tree does not exist -- no flake.nix, no ptyshim.c, no repo -- and
      # the image is already built, so validating and rebuilding it would
      # only fail. The App still has to exist for the decorators to bind.
      if self.is_remote:
          self.image = None
          self.app = modal.App(self.name)
      else:
          self._validate()
          self.image = self._build_image()
          self.app = modal.App(self.name, image=self.image)
  # -- spec reading ----------------------------------------------------
  @classmethod
  def from_toml(cls, container_py: str) -> "Container":
      """Load the container.toml sitting next to the given container.py."""
      directory = os.path.dirname(os.path.abspath(container_py))
      path = os.path.join(directory, "container.toml")
      if not os.path.exists(path):
          raise SpecError(f"no container.toml in {directory}")
      with open(path, "rb") as f:
          return cls(tomllib.load(f), directory)
  def _require(self, table: str, key: str):
      try:
          return self.spec[table][key]
      except KeyError:
          raise SpecError(f"container.toml needs [{table}] {key}") from None
  def _validate(self):
      c = self.spec.get("container", {})
      if self.runtime not in ("function", "sandbox"):
          raise SpecError(
              f'[container] runtime must be "function" or "sandbox",'
              f" not {self.runtime!r}"
      if bool(c.get("base")) == bool(c.get("registry")):
          raise SpecError("set exactly one of [container] base or registry")
      if self.use_flake:
          if not os.path.exists(os.path.join(self.dir, "flake.nix")):
              raise SpecError("[nix] flake = true but there is no flake.nix")
          if not self.use_shim:
              # Warming the shell builds nix-shell-env, which is the gVisor
              # pty bug. Failing here beats failing ten minutes into a build.
              raise SpecError(
                  "[nix] flake = true needs shim = true -- see ../README.md"
      if self.use_shim and not os.path.exists(PTYSHIM_C):
          raise SpecError(f"[nix] shim = true but {PTYSHIM_C} is missing")
A third target, and the seam that was already waiting for it   138        if self.ports and self.runtime != "sandbox":
          raise SpecError(
              "[network] ports needs [container] runtime = \"sandbox\" --"
              " a Function has no process to tunnel into"
Build the window somewhere with room for it   143        for name, mount in self.volume_spec.items():
          if not isinstance(mount, str) or not mount.startswith("/"):
              raise SpecError(
                  f"[volumes] {name} must be an absolute path, not {mount!r}"
          # Mounting over one of these hides what the image already has
          # there -- /nix in particular, where the empty volume would shadow
          # the store the base image spent its build populating.
          if mount.rstrip("/") in ("", "/nix", "/nix/store", "/usr", "/etc"):
              raise SpecError(
                  f"[volumes] {name} may not mount over {mount} --"
                  " it would hide what the image has there"
          if mount.rstrip("/") == self.workdir.rstrip("/"):
              raise SpecError(
                  f"[volumes] {name} may not mount over the workdir"
                  f" ({mount}) -- [build] include copies land there"
  # -- image -----------------------------------------------------------
  @property
  def nix(self) -> str:
      """The `nix` command as the container runs it.
      A sandbox on a real VM has a working pty, so the shim buys nothing
      there and is left off even when the image was built with it.
      """
      if self.use_shim and not self.vm_at_runtime:
          return f"LD_PRELOAD={SHIM_SO} nix"
      return "nix"
  @property
  def vm_at_runtime(self) -> bool:
      """True when this container actually runs on a VM rather than gVisor."""
      return (self.runtime == "sandbox"
              and bool(self.experimental_options.get("vm_runtime")))
  @property
  def build_nix(self) -> str:
      """The `nix` command for BUILD steps, which always run under gVisor.
      Image builds are Functions underneath, so a sandbox container still
      needs the shim while its image is being built -- only its run time
      gets the VM.
      """
      return f"LD_PRELOAD={SHIM_SO} nix" if self.use_shim else "nix"
  def _build_image(self) -> modal.Image:
      c = self.spec["container"]
      build = self.spec.get("build", {})
      if c.get("base"):
          image = modal.Image.from_name(c["base"])
      else:
          image = modal.Image.from_registry(c["registry"])
A shell that is already the devShell, and two recipes to reach it   200        # `nix profile install` puts things in ~/.nix-profile/bin, which is on
      # nobody's PATH here -- so the install succeeds and the very next line
      # says `command not found`. Set on the image rather than in a .bashrc,
      # because a container's command runs under `sh -c` and reads neither.
      # Spelled out rather than prefixed onto $PATH: an image env is a value,
      # not a shell expression, so "$PATH" here would be four literal
      # characters.
      image = image.env({
          "PATH": "/root/.nix-profile/bin:/usr/local/sbin:/usr/local/bin"
                  ":/usr/sbin:/usr/bin:/sbin:/bin",
Build the window somewhere with room for it   212        if self.use_shim:
          image = image.add_local_file(
              PTYSHIM_C, "/opt/ptyshim.c", copy=True
          ).run_commands(
              f"gcc -shared -fPIC -O2 -o {SHIM_SO} /opt/ptyshim.c -ldl"
      # Warm the devShell BEFORE the source is copied in. Everything the
      # shell needs is binary-cached, so the download happens once -- but
      # only if this layer survives. Copy the source first and any edit to
      # any file invalidates the warm, and the whole closure is fetched
      # again on every build. Only flake.nix and flake.lock go in here, so
      # the layer is invalidated by a dependency change and nothing else.
      if self.use_flake:
          for f in ("flake.nix", "flake.lock"):
              path = os.path.join(self.dir, f)
              if os.path.exists(path):
                  image = image.add_local_file(
                      path, f"{self.workdir}/{f}", copy=True
          image = image.run_commands(
              f"cd {self.workdir} && {self.build_nix} develop"
              " --accept-flake-config --command true",
              f'echo "store paths after warming:'
              f' $({self.build_nix} path-info --all | wc -l)"',
      # copy=True throughout: later run_commands need these files present.
      # `context` is what include paths are relative to, and it may sit above
      # the container directory -- a container that builds the repo it lives
      # in sets context = "../..", so include = ["."] means the whole repo.
      context = os.path.normpath(os.path.join(self.dir, build.get("context", ".")))
A third target, and the seam that was already waiting for it   244
      # Image building and program building, kept apart.
      # `[build] commands` run AFTER the source copy, so any edit anywhere in
      # the tree invalidates them -- and for a container whose commands warm
      # a devShell, that means minutes of nix on every iteration of a
      # one-line change. `warm` + `setup` are the same two steps moved in
      # front of the source: `warm` names only the files the flake actually
      # evaluates (its own, and whatever the devShell's derivations read),
      # and `setup` runs against those alone. Editing `flutter/src` then
      # invalidates nothing above the final copy, and the toolchain layer is
      # reused until a dependency moves.
      # Volumes are mounted here for `commands`' reason: a step that can
      # reach the binary cache never has to build, which is the one thing
      # gVisor will not do.
      for rel in build.get("warm", []):
          src = os.path.normpath(os.path.join(context, rel))
          dest = f"{self.workdir}/{rel}"
          if os.path.isdir(src):
              image = image.add_local_dir(src, dest, copy=True)
          else:
              image = image.add_local_file(src, dest, copy=True)
      if setup := build.get("setup", []):
          image = image.run_commands(*setup, volumes=self.volumes)
Stop shipping the build tree, and let a shell be the devShell   269        # `ignore` is what keeps a build tree out of the image. A checkout that
      # has been built in locally carries its output -- flutter/build and the
      # caches beside it were 395MB of a 441MB repo -- and all of it would be
      # uploaded on every start only to be thrown away, since the container
      # builds into a volume of its own. Patterns are relative to the copied
      # directory, as in .dockerignore.
      ignore = list(build.get("ignore", []))
Build the window somewhere with room for it   276        for rel in build.get("include", ["."]):
          src = os.path.normpath(os.path.join(context, rel))
          dest = self.workdir if rel == "." else f"{self.workdir}/{rel}"
          if os.path.isdir(src):
Stop shipping the build tree, and let a shell be the devShell   280                image = image.add_local_dir(src, dest, copy=True, ignore=ignore)
Build the window somewhere with room for it   281            else:
              image = image.add_local_file(src, dest, copy=True)
Carry the loader's .git fix over   284        # A repo copied in brings its `.git` along, and in a worktree that is a
      # *file* holding `gitdir: <path on the machine that copied it>`. Nix
      # believes it and goes looking for a checkout that is not there --
      # `nix develop` and `nix build .#x` both die before evaluating
      # anything. Nothing in a container wants the git metadata, so it goes.
      image = image.run_commands(f"rm -rf {self.workdir}/.git")
Build the window somewhere with room for it   291        if commands := build.get("commands", []):
A shell that is already the devShell, and two recipes to reach it   292            # Volumes mounted for the build too, not just the run. A build step
          # that wants nix to *build* something is the one thing gVisor will
          # not do -- image builds are Functions underneath, and a derivation
          # there dies on `unexpected EOF reading a line`. Substituting is
          # fine, so a step that can reach the cache never has to build, and
          # the shim stays retired. The mount is not part of the resulting
          # image; only what the step writes outside it is.
          image = image.run_commands(*commands, volumes=self.volumes)
Build the window somewhere with room for it   300
      # container.py does `from _loader import Container`, and Modal mounts
      # the entrypoint file alone -- so without this the import that works
      # locally fails in the container. copy=False adds it at startup rather
      # than baking a layer, so it invalidates nothing above it, and it must
      # therefore come after every build step.
      image = image.add_local_python_source("_loader")
      # ...and container.py reads its spec at import time, so the spec has to
      # be there too. Modal re-imports the entrypoint at /root, which is the
      # one directory that gets none of the workdir copies above.
      image = image.add_local_file(
          os.path.join(self.dir, "container.toml"), "/root/container.toml"
      return image
  # -- volumes ---------------------------------------------------------
  @property
  def volumes(self) -> dict:
      """{mount path: Volume}, as both Sandbox.create and @app.function want.
      `from_name` is lazy, so this is safe to evaluate on the re-import
      inside the container as well as out here. create_if_missing means a
      spec naming a volume that does not exist yet makes it rather than
      failing -- the first run of a cache is the one that fills it.
      """
      return {
          mount: modal.Volume.from_name(name, create_if_missing=True)
          for name, mount in self.volume_spec.items()
  # -- function --------------------------------------------------------
  @property
  def function_kwargs(self) -> dict:
      """Everything `@app.function` should be given, from [resources]."""
      r = self.spec.get("resources", {})
      kwargs: dict = {"timeout": int(r.get("timeout", 900))}
      if r.get("cpu"):
          kwargs["cpu"] = float(r["cpu"])
      if r.get("memory"):
          kwargs["memory"] = int(r["memory"])
      if r.get("gpu"):
          kwargs["gpu"] = r["gpu"]
      # Deliberately NOT self.experimental_options: that fills in the
      # sandbox default, and vm_runtime on a Function is refused by the
      # server. A sandbox container's @app.function is vestigial anyway.
      if self.runtime == "function":
          if experimental := dict(self.spec.get("experimental", {})):
              kwargs["experimental_options"] = experimental
      if volumes := self.volumes:
          kwargs["volumes"] = volumes
      return kwargs
  @property
  def experimental_options(self) -> dict:
      """Experimental options, with the sandbox default filled in.
      `vm_runtime` is Sandbox-only: the server rejects it on a Function
      outright. So it is defaulted on for sandboxes and never for functions,
      and an explicit [experimental] table always wins.
      """
      explicit = dict(self.spec.get("experimental", {}))
      if explicit:
          return explicit
      if self.runtime == "sandbox":
          return {"vm_runtime": True}
      return {}
  @property
  def sandbox_kwargs(self) -> dict:
      """Everything `Sandbox.create` should be given, from [resources]."""
      r = self.spec.get("resources", {})
      kwargs: dict = {"timeout": int(r.get("timeout", 900))}
      if r.get("cpu"):
          kwargs["cpu"] = float(r["cpu"])
      if r.get("memory"):
          # A VM sandbox gets exactly this much and cannot grow into more.
          kwargs["memory"] = int(r["memory"])
      if opts := self.experimental_options:
          kwargs["experimental_options"] = dict(opts)
      if volumes := self.volumes:
          kwargs["volumes"] = volumes
A third target, and the seam that was already waiting for it   384        if self.ports:
          kwargs["encrypted_ports"] = list(self.ports)
Build the window somewhere with room for it   386        return kwargs
  def shell_command(self, override: str = "") -> str:
      """The full shell line the container runs, devShell wrapper included."""
      command = override or self.command
      if not command:
          raise SpecError("container.toml has no [run] command")
      if self.use_flake:
          return (
              f"cd {self.workdir} && {self.nix} develop --accept-flake-config"
              f" --command sh -c {shlex.quote(command)}"
      return f"cd {self.workdir} && {command}"
  def run_sandbox(self, override: str = "") -> str:
      """Run one command in a Sandbox that dies when the command does.
      The command IS the sandbox's process, rather than something exec'd
      into a `sleep infinity` box that then has to be torn down. There is no
      idle window to pay for and nothing to leak if this script is killed;
      the [resources] timeout is a backstop, not the mechanism.
      Modal streams a Sandbox's output into the app log as it runs -- which
      is what you want for a build -- so this returns "" rather than handing
      back a copy for the caller to print underneath it.
      """
      line = self.shell_command(override)
      sb = modal.Sandbox.create(
          "sh", "-c", line,
          app=self.app,
          image=self.image,
          workdir=self.workdir,
          env={k: str(v) for k, v in self.env.items()},
          **self.sandbox_kwargs,
A third target, and the seam that was already waiting for it   421        self._print_tunnels(sb)
Build the window somewhere with room for it   422        sb.wait()
      if sb.returncode != 0:
          raise RuntimeError(
              f"{self.name}: command failed ({sb.returncode})\n"
              f"$ {line}\n{sb.stderr.read()}"
      return ""
A third target, and the seam that was already waiting for it   430    def _print_tunnels(self, sb: "modal.Sandbox") -> None:
      """Say where a tunnelled port can be reached, once the sandbox is up.
      `tunnels()` blocks until the Sandbox is scheduled, which is why this
      is called after create and not folded into it. Nothing prints when
      [network] ports is empty, which is every container but the serving
      ones.
      """
      if not self.ports:
          return
      for port, tunnel in sb.tunnels().items():
          print(f"  :{port} -> {tunnel.url}")
Build the window somewhere with room for it   443    def open_sandbox(self) -> "modal.Sandbox":
Stop shipping the build tree, and let a shell be the devShell   444        """Start a Sandbox and leave it running, for `scripts/shell`.
      Same workdir and the same [run] env as the real thing: a shell opened
      to debug a container that does not have the container's environment is
      a shell that reproduces something else. `command` is the one part left
      out, because not running it is the point.
      """
Build the window somewhere with room for it   451        return modal.Sandbox.create(
          "sleep",
          "infinity",
          app=self.app,
          image=self.image,
Stop shipping the build tree, and let a shell be the devShell   456            workdir=self.workdir,
          env={k: str(v) for k, v in self.env.items()},
Build the window somewhere with room for it   458            **self.sandbox_kwargs,
  def execute(self, override: str = "") -> str:
      """Run the command in this container. Called remotely, not locally."""
      line = self.shell_command(override)
      result = subprocess.run(
          line,
          shell=True,
          capture_output=True,
          text=True,
          env={**os.environ, **{k: str(v) for k, v in self.env.items()}},
      if result.returncode != 0:
          raise RuntimeError(
              f"{self.name}: command failed ({result.returncode})\n"
              f"$ {line}\n{result.stderr}"
      return result.stdout