| Build the window somewhere with room for it 1451151 nandi 5d ago | 1 | """Turn a `container.toml` into a `modal.Image` and a `modal.App`. |
| 2 | |
| 3 | Every container under `containers/` is a directory with a `container.toml` and |
| 4 | a stub `container.py`. See `spec.md` for the keys; this file is what reads |
| 5 | them. Nothing here is Modal-specific configuration in its own right -- each |
| 6 | spec key maps onto a documented Modal argument, and the mapping is meant to |
| 7 | stay boring enough to read straight through. |
| 8 | """ |
| 9 | |
| 10 | import os |
| 11 | import shlex |
| 12 | import subprocess |
| 13 | import tomllib |
| 14 | |
| 15 | import modal |
| 16 | |
| 17 | REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| 18 | PTYSHIM_C = os.path.join(REPO, "ptyshim.c") |
| 19 | SHIM_SO = "/opt/ptyshim.so" |
| 20 | |
| 21 | |
| 22 | class SpecError(Exception): |
| 23 | """The container.toml says something that cannot be built.""" |
| 24 | |
| 25 | |
| 26 | def _running_in_modal() -> bool: |
| 27 | """True inside a Modal container, false on the machine that launched it.""" |
| 28 | return bool(os.environ.get("MODAL_TASK_ID")) |
| 29 | |
| 30 | |
| 31 | class Container: |
| 32 | """One container: its spec, its image, its app, and how to run it.""" |
| 33 | |
| 34 | def __init__(self, spec: dict, directory: str): |
| 35 | self.is_remote = _running_in_modal() |
| 36 | self.dir = directory |
| 37 | self.spec = spec |
| 38 | self.name = self._require("container", "name") |
| 39 | self.description = spec.get("container", {}).get("description", "") |
| 40 | |
| 41 | run = spec.get("run", {}) |
| 42 | self.workdir = run.get("workdir", "/app") |
| 43 | self.command = run.get("command", "") |
| 44 | self.env = dict(run.get("env", {})) |
| 45 | |
| 46 | # "function" (the default) or "sandbox". Sandboxes can run on a real |
| 47 | # VM, which Functions cannot -- see ../README.md. |
| 48 | self.runtime = spec.get("container", {}).get("runtime", "function") |
| 49 | |
| 50 | nix = spec.get("nix", {}) |
| 51 | self.use_flake = bool(nix.get("flake", False)) |
| 52 | self.use_shim = bool(nix.get("shim", False)) |
| 53 | |
| 54 | # name -> mount path. Modal Volumes, mounted while the container runs |
| 55 | # and NOT while its image is built: a volume mount is not part of the |
| 56 | # resulting image, so anything written to one during a build step is |
| 57 | # gone by the time the container starts. Persist across runs is the |
| 58 | # whole point -- a nix store to substitute from, a cargo target |
| 59 | # directory, a dataset too big to bake in. |
| 60 | self.volume_spec = dict(spec.get("volumes", {})) |
| 61 | |
| 62 | # Modal re-imports this module inside the container, so everything |
| 63 | # below runs twice: once here, once out there. Out there the local |
| 64 | # tree does not exist -- no flake.nix, no ptyshim.c, no repo -- and |
| 65 | # the image is already built, so validating and rebuilding it would |
| 66 | # only fail. The App still has to exist for the decorators to bind. |
| 67 | if self.is_remote: |
| 68 | self.image = None |
| 69 | self.app = modal.App(self.name) |
| 70 | else: |
| 71 | self._validate() |
| 72 | self.image = self._build_image() |
| 73 | self.app = modal.App(self.name, image=self.image) |
| 74 | |
| 75 | # -- spec reading ---------------------------------------------------- |
| 76 | |
| 77 | @classmethod |
| 78 | def from_toml(cls, container_py: str) -> "Container": |
| 79 | """Load the container.toml sitting next to the given container.py.""" |
| 80 | directory = os.path.dirname(os.path.abspath(container_py)) |
| 81 | path = os.path.join(directory, "container.toml") |
| 82 | if not os.path.exists(path): |
| 83 | raise SpecError(f"no container.toml in {directory}") |
| 84 | with open(path, "rb") as f: |
| 85 | return cls(tomllib.load(f), directory) |
| 86 | |
| 87 | def _require(self, table: str, key: str): |
| 88 | try: |
| 89 | return self.spec[table][key] |
| 90 | except KeyError: |
| 91 | raise SpecError(f"container.toml needs [{table}] {key}") from None |
| 92 | |
| 93 | def _validate(self): |
| 94 | c = self.spec.get("container", {}) |
| 95 | if self.runtime not in ("function", "sandbox"): |
| 96 | raise SpecError( |
| 97 | f'[container] runtime must be "function" or "sandbox",' |
| 98 | f" not {self.runtime!r}" |
| 99 | ) |
| 100 | if bool(c.get("base")) == bool(c.get("registry")): |
| 101 | raise SpecError("set exactly one of [container] base or registry") |
| 102 | if self.use_flake: |
| 103 | if not os.path.exists(os.path.join(self.dir, "flake.nix")): |
| 104 | raise SpecError("[nix] flake = true but there is no flake.nix") |
| 105 | if not self.use_shim: |
| 106 | # Warming the shell builds nix-shell-env, which is the gVisor |
| 107 | # pty bug. Failing here beats failing ten minutes into a build. |
| 108 | raise SpecError( |
| 109 | "[nix] flake = true needs shim = true -- see ../README.md" |
| 110 | ) |
| 111 | if self.use_shim and not os.path.exists(PTYSHIM_C): |
| 112 | raise SpecError(f"[nix] shim = true but {PTYSHIM_C} is missing") |
| 113 | for name, mount in self.volume_spec.items(): |
| 114 | if not isinstance(mount, str) or not mount.startswith("/"): |
| 115 | raise SpecError( |
| 116 | f"[volumes] {name} must be an absolute path, not {mount!r}" |
| 117 | ) |
| 118 | # Mounting over one of these hides what the image already has |
| 119 | # there -- /nix in particular, where the empty volume would shadow |
| 120 | # the store the base image spent its build populating. |
| 121 | if mount.rstrip("/") in ("", "/nix", "/nix/store", "/usr", "/etc"): |
| 122 | raise SpecError( |
| 123 | f"[volumes] {name} may not mount over {mount} --" |
| 124 | " it would hide what the image has there" |
| 125 | ) |
| 126 | if mount.rstrip("/") == self.workdir.rstrip("/"): |
| 127 | raise SpecError( |
| 128 | f"[volumes] {name} may not mount over the workdir" |
| 129 | f" ({mount}) -- [build] include copies land there" |
| 130 | ) |
| 131 | |
| 132 | # -- image ----------------------------------------------------------- |
| 133 | |
| 134 | @property |
| 135 | def nix(self) -> str: |
| 136 | """The `nix` command as the container runs it. |
| 137 | |
| 138 | A sandbox on a real VM has a working pty, so the shim buys nothing |
| 139 | there and is left off even when the image was built with it. |
| 140 | """ |
| 141 | if self.use_shim and not self.vm_at_runtime: |
| 142 | return f"LD_PRELOAD={SHIM_SO} nix" |
| 143 | return "nix" |
| 144 | |
| 145 | @property |
| 146 | def vm_at_runtime(self) -> bool: |
| 147 | """True when this container actually runs on a VM rather than gVisor.""" |
| 148 | return (self.runtime == "sandbox" |
| 149 | and bool(self.experimental_options.get("vm_runtime"))) |
| 150 | |
| 151 | @property |
| 152 | def build_nix(self) -> str: |
| 153 | """The `nix` command for BUILD steps, which always run under gVisor. |
| 154 | |
| 155 | Image builds are Functions underneath, so a sandbox container still |
| 156 | needs the shim while its image is being built -- only its run time |
| 157 | gets the VM. |
| 158 | """ |
| 159 | return f"LD_PRELOAD={SHIM_SO} nix" if self.use_shim else "nix" |
| 160 | |
| 161 | def _build_image(self) -> modal.Image: |
| 162 | c = self.spec["container"] |
| 163 | build = self.spec.get("build", {}) |
| 164 | |
| 165 | if c.get("base"): |
| 166 | image = modal.Image.from_name(c["base"]) |
| 167 | else: |
| 168 | image = modal.Image.from_registry(c["registry"]) |
| 169 | |
| 170 | if self.use_shim: |
| 171 | image = image.add_local_file( |
| 172 | PTYSHIM_C, "/opt/ptyshim.c", copy=True |
| 173 | ).run_commands( |
| 174 | f"gcc -shared -fPIC -O2 -o {SHIM_SO} /opt/ptyshim.c -ldl" |
| 175 | ) |
| 176 | |
| 177 | # Warm the devShell BEFORE the source is copied in. Everything the |
| 178 | # shell needs is binary-cached, so the download happens once -- but |
| 179 | # only if this layer survives. Copy the source first and any edit to |
| 180 | # any file invalidates the warm, and the whole closure is fetched |
| 181 | # again on every build. Only flake.nix and flake.lock go in here, so |
| 182 | # the layer is invalidated by a dependency change and nothing else. |
| 183 | if self.use_flake: |
| 184 | for f in ("flake.nix", "flake.lock"): |
| 185 | path = os.path.join(self.dir, f) |
| 186 | if os.path.exists(path): |
| 187 | image = image.add_local_file( |
| 188 | path, f"{self.workdir}/{f}", copy=True |
| 189 | ) |
| 190 | image = image.run_commands( |
| 191 | f"cd {self.workdir} && {self.build_nix} develop" |
| 192 | " --accept-flake-config --command true", |
| 193 | f'echo "store paths after warming:' |
| 194 | f' $({self.build_nix} path-info --all | wc -l)"', |
| 195 | ) |
| 196 | |
| 197 | # copy=True throughout: later run_commands need these files present. |
| 198 | # `context` is what include paths are relative to, and it may sit above |
| 199 | # the container directory -- a container that builds the repo it lives |
| 200 | # in sets context = "../..", so include = ["."] means the whole repo. |
| 201 | context = os.path.normpath(os.path.join(self.dir, build.get("context", "."))) |
| 202 | for rel in build.get("include", ["."]): |
| 203 | src = os.path.normpath(os.path.join(context, rel)) |
| 204 | dest = self.workdir if rel == "." else f"{self.workdir}/{rel}" |
| 205 | if os.path.isdir(src): |
| 206 | image = image.add_local_dir(src, dest, copy=True) |
| 207 | else: |
| 208 | image = image.add_local_file(src, dest, copy=True) |
| 209 | |
| 210 | if commands := build.get("commands", []): |
| 211 | image = image.run_commands(*commands) |
| 212 | |
| 213 | # container.py does `from _loader import Container`, and Modal mounts |
| 214 | # the entrypoint file alone -- so without this the import that works |
| 215 | # locally fails in the container. copy=False adds it at startup rather |
| 216 | # than baking a layer, so it invalidates nothing above it, and it must |
| 217 | # therefore come after every build step. |
| 218 | image = image.add_local_python_source("_loader") |
| 219 | # ...and container.py reads its spec at import time, so the spec has to |
| 220 | # be there too. Modal re-imports the entrypoint at /root, which is the |
| 221 | # one directory that gets none of the workdir copies above. |
| 222 | image = image.add_local_file( |
| 223 | os.path.join(self.dir, "container.toml"), "/root/container.toml" |
| 224 | ) |
| 225 | |
| 226 | return image |
| 227 | |
| 228 | # -- volumes --------------------------------------------------------- |
| 229 | |
| 230 | @property |
| 231 | def volumes(self) -> dict: |
| 232 | """{mount path: Volume}, as both Sandbox.create and @app.function want. |
| 233 | |
| 234 | `from_name` is lazy, so this is safe to evaluate on the re-import |
| 235 | inside the container as well as out here. create_if_missing means a |
| 236 | spec naming a volume that does not exist yet makes it rather than |
| 237 | failing -- the first run of a cache is the one that fills it. |
| 238 | """ |
| 239 | return { |
| 240 | mount: modal.Volume.from_name(name, create_if_missing=True) |
| 241 | for name, mount in self.volume_spec.items() |
| 242 | } |
| 243 | |
| 244 | # -- function -------------------------------------------------------- |
| 245 | |
| 246 | @property |
| 247 | def function_kwargs(self) -> dict: |
| 248 | """Everything `@app.function` should be given, from [resources].""" |
| 249 | r = self.spec.get("resources", {}) |
| 250 | kwargs: dict = {"timeout": int(r.get("timeout", 900))} |
| 251 | if r.get("cpu"): |
| 252 | kwargs["cpu"] = float(r["cpu"]) |
| 253 | if r.get("memory"): |
| 254 | kwargs["memory"] = int(r["memory"]) |
| 255 | if r.get("gpu"): |
| 256 | kwargs["gpu"] = r["gpu"] |
| 257 | # Deliberately NOT self.experimental_options: that fills in the |
| 258 | # sandbox default, and vm_runtime on a Function is refused by the |
| 259 | # server. A sandbox container's @app.function is vestigial anyway. |
| 260 | if self.runtime == "function": |
| 261 | if experimental := dict(self.spec.get("experimental", {})): |
| 262 | kwargs["experimental_options"] = experimental |
| 263 | if volumes := self.volumes: |
| 264 | kwargs["volumes"] = volumes |
| 265 | return kwargs |
| 266 | |
| 267 | @property |
| 268 | def experimental_options(self) -> dict: |
| 269 | """Experimental options, with the sandbox default filled in. |
| 270 | |
| 271 | `vm_runtime` is Sandbox-only: the server rejects it on a Function |
| 272 | outright. So it is defaulted on for sandboxes and never for functions, |
| 273 | and an explicit [experimental] table always wins. |
| 274 | """ |
| 275 | explicit = dict(self.spec.get("experimental", {})) |
| 276 | if explicit: |
| 277 | return explicit |
| 278 | if self.runtime == "sandbox": |
| 279 | return {"vm_runtime": True} |
| 280 | return {} |
| 281 | |
| 282 | @property |
| 283 | def sandbox_kwargs(self) -> dict: |
| 284 | """Everything `Sandbox.create` should be given, from [resources].""" |
| 285 | r = self.spec.get("resources", {}) |
| 286 | kwargs: dict = {"timeout": int(r.get("timeout", 900))} |
| 287 | if r.get("cpu"): |
| 288 | kwargs["cpu"] = float(r["cpu"]) |
| 289 | if r.get("memory"): |
| 290 | # A VM sandbox gets exactly this much and cannot grow into more. |
| 291 | kwargs["memory"] = int(r["memory"]) |
| 292 | if opts := self.experimental_options: |
| 293 | kwargs["experimental_options"] = dict(opts) |
| 294 | if volumes := self.volumes: |
| 295 | kwargs["volumes"] = volumes |
| 296 | return kwargs |
| 297 | |
| 298 | def shell_command(self, override: str = "") -> str: |
| 299 | """The full shell line the container runs, devShell wrapper included.""" |
| 300 | command = override or self.command |
| 301 | if not command: |
| 302 | raise SpecError("container.toml has no [run] command") |
| 303 | if self.use_flake: |
| 304 | return ( |
| 305 | f"cd {self.workdir} && {self.nix} develop --accept-flake-config" |
| 306 | f" --command sh -c {shlex.quote(command)}" |
| 307 | ) |
| 308 | return f"cd {self.workdir} && {command}" |
| 309 | |
| 310 | def run_sandbox(self, override: str = "") -> str: |
| 311 | """Run one command in a Sandbox that dies when the command does. |
| 312 | |
| 313 | The command IS the sandbox's process, rather than something exec'd |
| 314 | into a `sleep infinity` box that then has to be torn down. There is no |
| 315 | idle window to pay for and nothing to leak if this script is killed; |
| 316 | the [resources] timeout is a backstop, not the mechanism. |
| 317 | |
| 318 | Modal streams a Sandbox's output into the app log as it runs -- which |
| 319 | is what you want for a build -- so this returns "" rather than handing |
| 320 | back a copy for the caller to print underneath it. |
| 321 | """ |
| 322 | line = self.shell_command(override) |
| 323 | sb = modal.Sandbox.create( |
| 324 | "sh", "-c", line, |
| 325 | app=self.app, |
| 326 | image=self.image, |
| 327 | workdir=self.workdir, |
| 328 | env={k: str(v) for k, v in self.env.items()}, |
| 329 | **self.sandbox_kwargs, |
| 330 | ) |
| 331 | sb.wait() |
| 332 | if sb.returncode != 0: |
| 333 | raise RuntimeError( |
| 334 | f"{self.name}: command failed ({sb.returncode})\n" |
| 335 | f"$ {line}\n{sb.stderr.read()}" |
| 336 | ) |
| 337 | return "" |
| 338 | |
| 339 | def open_sandbox(self) -> "modal.Sandbox": |
| 340 | """Start a Sandbox and leave it running, for `scripts/shell`.""" |
| 341 | return modal.Sandbox.create( |
| 342 | "sleep", |
| 343 | "infinity", |
| 344 | app=self.app, |
| 345 | image=self.image, |
| 346 | **self.sandbox_kwargs, |
| 347 | ) |
| 348 | |
| 349 | def execute(self, override: str = "") -> str: |
| 350 | """Run the command in this container. Called remotely, not locally.""" |
| 351 | line = self.shell_command(override) |
| 352 | result = subprocess.run( |
| 353 | line, |
| 354 | shell=True, |
| 355 | capture_output=True, |
| 356 | text=True, |
| 357 | env={**os.environ, **{k: str(v) for k, v in self.env.items()}}, |
| 358 | ) |
| 359 | if result.returncode != 0: |
| 360 | raise RuntimeError( |
| 361 | f"{self.name}: command failed ({result.returncode})\n" |
| 362 | f"$ {line}\n{result.stderr}" |
| 363 | ) |
| 364 | return result.stdout |