dimos bake: compose rust native modules into one host binary

dimos/cli/bake/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 2026 Dimensional Inc.
+#
+# 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.
+
+"""`dimos bake` — compose rust native modules into one host binary."""
+
+from __future__ import annotations
+
+
+class BakeError(Exception):
+    """A bake that cannot proceed: bad registry, unwireable graph, failed build."""

dimos/cli/bake/build.py

@@ -0,0 +1,84 @@
+# Copyright 2026 Dimensional Inc.
+#
+# 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.
+
+"""Invoke a cargo-shaped builder on the generated crate and collect the binary.
+
+Deliberately a thin invoker: `cross` and `cargo-zigbuild` take the same
+arguments as cargo, and managing their toolchains is not bake's job.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+import shutil
+import subprocess
+
+from dimos.cli.bake import BakeError
+
+BUILDERS = ("cargo", "cross", "zigbuild")
+
+_INVOCATION = {
+    "cargo": ["cargo", "build"],
+    "cross": ["cross", "build"],
+    "zigbuild": ["cargo", "zigbuild"],
+}
+
+
+def build_command(builder: str, *, target: str | None = None, debug: bool = False) -> list[str]:
+    if builder not in _INVOCATION:
+        raise BakeError(f"unknown --builder {builder!r}; choose from {', '.join(BUILDERS)}")
+    cmd = list(_INVOCATION[builder])
+    if not debug:
+        cmd.append("--release")
+    if target:
+        cmd.extend(["--target", target])
+    return cmd
+
+
+def artifact_path(
+    crate_dir: Path, host: str, *, target: str | None = None, debug: bool = False
+) -> Path:
+    profile = "debug" if debug else "release"
+    out = crate_dir / "target"
+    if target:
+        out = out / target
+    return out / profile / host
+
+
+def build_host(
+    crate_dir: Path,
+    host: str,
+    *,
+    builder: str = "cargo",
+    target: str | None = None,
+    debug: bool = False,
+) -> Path:
+    """Compile the generated crate and return the path to the built binary."""
+    cmd = build_command(builder, target=target, debug=debug)
+    if shutil.which(cmd[0]) is None:
+        raise BakeError(f"`{cmd[0]}` is not on PATH")
+    result = subprocess.run(cmd, cwd=crate_dir, check=False)
+    if result.returncode != 0:
+        raise BakeError(f"{' '.join(cmd)} failed with exit {result.returncode}")
+    artifact = artifact_path(crate_dir, host, target=target, debug=debug)
+    if not artifact.exists():
+        raise BakeError(f"build succeeded but {artifact} is missing")
+    return artifact
+
+
+def install(artifact: Path, out: Path) -> int:
+    """Copy the built binary to `out`, returning its size in bytes."""
+    out.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(artifact, out)
+    return out.stat().st_size

dimos/cli/bake/cli.py

@@ -0,0 +1,149 @@
+# Copyright 2026 Dimensional Inc.
+#
+# 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.
+
+"""`dimos bake` — link native modules into one host binary."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import importlib
+import json
+from pathlib import Path
+from typing import Any, get_type_hints
+
+import typer
+
+from dimos.cli.bake import BakeError
+from dimos.cli.bake.build import BUILDERS, build_host, install
+from dimos.cli.bake.codegen import generate_crate
+from dimos.cli.bake.discovery import ModuleInfo, discover_modules, render_registry, select_modules
+from dimos.cli.bake.graph import Graph, build_graph, parse_remap, render
+
+
+def default_config(module: ModuleInfo) -> dict[str, Any]:
+    """The module's python wrapper config at its defaults, as the native struct sees it."""
+    module_name, _, class_name = module.python_ref.partition(":")
+    try:
+        wrapper = getattr(importlib.import_module(module_name), class_name)
+    except (ImportError, AttributeError) as exc:
+        raise BakeError(
+            f"module `{module.id}`: cannot import `{module.python_ref}`: {exc}"
+        ) from exc
+    config_type = get_type_hints(wrapper)["config"]
+    return dict(config_type().to_config_dict())
+
+
+def emit_config(graph: Graph, modules: Sequence[ModuleInfo]) -> dict[str, Any]:
+    """A complete stdin blob for the host, so it can be run without python."""
+    topics = graph.topics()
+    return {
+        "modules": {
+            module.id: {
+                "topics": topics[module.id],
+                "config": default_config(module) or None,
+            }
+            for module in modules
+        },
+        "qos": graph.qos(),
+        "suppress": list(graph.suppressed_topics()),
+    }
+
+
+def bake(
+    modules: list[str] = typer.Argument(
+        None, help="Module ids to bake, e.g. `ray-tracing mls-planner`."
+    ),
+    out: Path = typer.Option(
+        None, "-o", "--out", help="Where to write the host binary; its name names the host."
+    ),
+    target: str = typer.Option(None, "--target", help="Rust target triple. Default: host-native."),
+    suppress: list[str] = typer.Option(
+        None,
+        "--suppress",
+        help="Keep this topic inside the host (name or full topic). Repeatable.",
+    ),
+    remap: list[str] = typer.Option(
+        None, "--remap", help="Rename a port: <module>.<port>=<name>. Repeatable."
+    ),
+    builder: str = typer.Option("cargo", "--builder", help=f"Build driver: {', '.join(BUILDERS)}."),
+    debug: bool = typer.Option(False, "--debug", help="Build the dev profile instead of release."),
+    dry_run: bool = typer.Option(False, "--dry-run", help="Print the graph and stop."),
+    emit_config_to: Path = typer.Option(
+        None, "--emit-config", help="Also write a ready-to-pipe stdin JSON config here."
+    ),
+    list_modules: bool = typer.Option(
+        False, "--list", help="List registered native modules and exit."
+    ),
+) -> None:
+    """Compose rust native modules into a single host binary."""
+    try:
+        _bake(
+            modules or [],
+            out=out,
+            target=target,
+            suppress=suppress or [],
+            remap=remap or [],
+            builder=builder,
+            debug=debug,
+            dry_run=dry_run,
+            emit_config_to=emit_config_to,
+            list_modules=list_modules,
+        )
+    except BakeError as exc:
+        typer.echo(typer.style(f"bake: {exc}", fg=typer.colors.RED), err=True)
+        raise typer.Exit(1) from exc
+
+
+def _bake(
+    module_names: Sequence[str],
+    *,
+    out: Path | None,
+    target: str | None,
+    suppress: Sequence[str],
+    remap: Sequence[str],
+    builder: str,
+    debug: bool,
+    dry_run: bool,
+    emit_config_to: Path | None,
+    list_modules: bool,
+) -> None:
+    registry = discover_modules()
+    if list_modules:
+        typer.echo(render_registry(registry))
+        return
+
+    if out is None:
+        raise BakeError("-o/--out is required: its filename names the host binary")
+    host = out.name
+    selected = select_modules(registry, module_names)
+    graph = build_graph(
+        host, selected, remaps=dict(parse_remap(r) for r in remap), suppress=suppress
+    )
+
+    typer.echo(render(graph))
+    typer.echo("")
+
+    if emit_config_to is not None:
+        emit_config_to.parent.mkdir(parents=True, exist_ok=True)
+        emit_config_to.write_text(json.dumps(emit_config(graph, selected), indent=2) + "\n")
+        typer.echo(f"Wrote {emit_config_to}")
+
+    if dry_run:
+        return
+
+    crate = generate_crate(host, selected, graph)
+    typer.echo(f"Generated {crate}")
+    artifact = build_host(crate, host, builder=builder, target=target, debug=debug)
+    size = install(artifact, out)
+    typer.echo(f"Wrote {out} ({size / 1e6:.1f} MB)")

dimos/cli/bake/codegen.py

@@ -0,0 +1,141 @@
+# Copyright 2026 Dimensional Inc.
+#
+# 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.
+
+"""Emit the throwaway cargo crate that becomes the host binary.
+
+The crate is a table of module entries plus three JSON blobs; all the runtime
+behaviour lives in `dimos_module::host`, so there is nothing here to debug.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import json
+from pathlib import Path
+
+from dimos.cli.bake.discovery import ModuleInfo, repo_root
+from dimos.cli.bake.graph import Graph
+
+GENERATED_HEADER = "// GENERATED BY dimos bake — do not edit."
+
+_CARGO_TEMPLATE = """# GENERATED BY dimos bake — do not edit.
+[workspace]
+
+[package]
+name = "{host}"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[[bin]]
+name = "{host}"
+path = "src/main.rs"
+
+[dependencies]
+{dependencies}
+
+[profile.release]
+lto = "thin"
+codegen-units = 1
+strip = "symbols"
+"""
+
+_MAIN_TEMPLATE = """{header}
+use dimos_module::{{host_main, HostSpec, ModuleEntry}};
+
+static MODULES: &[ModuleEntry] = &[
+{entries}
+];
+
+static SUPPRESS: &[&str] = &[
+{suppress}
+];
+
+static SPEC: HostSpec = HostSpec {{
+    name: "{host}",
+    modules: MODULES,
+    default_topics: include_str!("default_topics.json"),
+    default_suppress: SUPPRESS,
+    default_qos: include_str!("default_qos.json"),
+    graph_json: include_str!("graph.json"),
+}};
+
+fn main() {{
+    host_main(&SPEC)
+}}
+"""
+
+
+def crate_dir(host: str, root: Path | None = None) -> Path:
+    """Where the generated crate for `host` lives."""
+    return (root or repo_root()) / "target" / "dimos-bake" / host
+
+
+def _dependencies(modules: Sequence[ModuleInfo], root: Path) -> str:
+    lines = [f'dimos-module = {{ path = "{root / "native" / "rust" / "dimos-module"}" }}']
+    for module in modules:
+        # default-features off drops pyo3, whose libpython symbols would break a
+        # static link and are dead weight in a host either way.
+        lines.append(
+            f'{module.crate_name} = {{ path = "{module.crate_dir}", default-features = false }}'
+        )
+    return "\n".join(lines)
+
+
+def _entries(modules: Sequence[ModuleInfo]) -> str:
+    lines = []
+    for module in modules:
+        call = f'    ModuleEntry::new::<{module.rust_path}>("{module.id}")'
+        if module.threads != 1:
+            call += f".threads({module.threads})"
+        if module.nice is not None:
+            call += f".nice({module.nice})"
+        lines.append(call + ",")
+    return "\n".join(lines)
+
+
+def render_cargo_toml(host: str, modules: Sequence[ModuleInfo], root: Path | None = None) -> str:
+    return _CARGO_TEMPLATE.format(
+        host=host, dependencies=_dependencies(modules, root or repo_root())
+    )
+
+
+def render_main_rs(host: str, modules: Sequence[ModuleInfo], graph: Graph) -> str:
+    suppress = "\n".join(f'    "{topic}",' for topic in graph.suppressed_topics())
+    return _MAIN_TEMPLATE.format(
+        header=GENERATED_HEADER,
+        host=host,
+        entries=_entries(modules),
+        suppress=suppress,
+    )
+
+
+def generate_crate(
+    host: str,
+    modules: Sequence[ModuleInfo],
+    graph: Graph,
+    root: Path | None = None,
+) -> Path:
+    """Write the generated crate and return its directory."""
+    root = root or repo_root()
+    directory = crate_dir(host, root)
+    src = directory / "src"
+    src.mkdir(parents=True, exist_ok=True)
+
+    (directory / "Cargo.toml").write_text(render_cargo_toml(host, modules, root))
+    (src / "main.rs").write_text(render_main_rs(host, modules, graph))
+    (src / "default_topics.json").write_text(json.dumps(graph.topics(), indent=2) + "\n")
+    (src / "default_qos.json").write_text(json.dumps(graph.qos(), indent=2) + "\n")
+    (src / "graph.json").write_text(json.dumps(graph.to_json(), indent=2) + "\n")
+    return directory