diff --git a/MODULE.bazel b/MODULE.bazel index ef6a202e3dfb..37de3f036cdd 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -306,6 +306,17 @@ http_archive( url = "https://github.com/apalache-mc/apalache/releases/download/v0.52.2/apalache-0.52.2.tgz", ) +# Ollama CPU-only runtime, shipped in the GuestOS. Only the `ollama` binary +# and the CPU ggml shared libraries are kept; GPU runners (CUDA/MLX/Vulkan) +# are stripped via the `build_file` filter. See BUILD.ollama.bazel. +http_archive( + name = "ollama", + build_file = "@//third_party:BUILD.ollama.bazel", + sha256 = "528acc42750f755ce57b13b8138419ceac953b6fd488045f73dc1e80246da480", + type = "tar.zst", + url = "https://github.com/ollama/ollama/releases/download/v0.21.1/ollama-linux-amd64.tar.zst", +) + # Official WebAssembly test suite. # To be used for testing libraries that handle canister Wasm code. http_archive( diff --git a/ic-os/components/guestos.bzl b/ic-os/components/guestos.bzl index 0732981f3917..17fcfa6d4315 100644 --- a/ic-os/components/guestos.bzl +++ b/ic-os/components/guestos.bzl @@ -50,6 +50,13 @@ def component_files(mode): Label("guestos/ic-https-outcalls-adapter/ic-https-outcalls-adapter.socket"): "/etc/systemd/system/ic-https-outcalls-adapter.socket", Label("guestos/ic-https-outcalls-adapter/generate-https-outcalls-adapter-config.sh"): "/opt/ic/bin/generate-https-outcalls-adapter-config.sh", Label("guestos/ic-replica.service"): "/etc/systemd/system/ic-replica.service", + # Ollama is intentionally disabled by default. The blanket + # `systemctl enable` loop in the Dockerfile would otherwise enable + # it; an explicit `systemctl disable ollama.service` in the + # Dockerfile keeps it off. Start it on demand with: + # systemctl enable --now ollama.service + Label("guestos/ollama/ollama.service"): "/etc/systemd/system/ollama.service", + Label("guestos/ollama/manage-ollama.sh"): "/opt/ic/bin/manage-ollama.sh", Label("guestos/remote-attestation-server.service"): "/etc/systemd/system/remote-attestation-server.service", Label("guestos/generate-ic-config/generate-ic-config.service"): "/etc/systemd/system/generate-ic-config.service", Label("guestos/share/ic-boundary.env"): "/opt/ic/share/ic-boundary.env", diff --git a/ic-os/components/guestos/misc/sudoers b/ic-os/components/guestos/misc/sudoers index 3adcb37a6bf4..b0f381b98ad7 100644 --- a/ic-os/components/guestos/misc/sudoers +++ b/ic-os/components/guestos/misc/sudoers @@ -29,6 +29,6 @@ root ALL=(ALL:ALL) NOPASSWD:ALL # Allow members of group sudo to execute any command %sudo ALL=(ALL:ALL) NOPASSWD:ALL -ic-replica ALL=(ALL:ALL) NOPASSWD: /opt/ic/bin/manageboot.sh, /opt/ic/bin/provision-ssh-keys.sh, /opt/ic/bin/read-ssh-keys.sh, /opt/ic/bin/sync_fstrim.sh, /opt/ic/bin/guestos_tool, /usr/sbin/nft +ic-replica ALL=(ALL:ALL) NOPASSWD: /opt/ic/bin/manageboot.sh, /opt/ic/bin/provision-ssh-keys.sh, /opt/ic/bin/read-ssh-keys.sh, /opt/ic/bin/sync_fstrim.sh, /opt/ic/bin/guestos_tool, /opt/ic/bin/manage-ollama.sh, /usr/sbin/nft # See sudoers(5) for more information on "#include" directives: diff --git a/ic-os/components/guestos/ollama/manage-ollama.sh b/ic-os/components/guestos/ollama/manage-ollama.sh new file mode 100755 index 000000000000..f56395137ce8 --- /dev/null +++ b/ic-os/components/guestos/ollama/manage-ollama.sh @@ -0,0 +1,24 @@ +#!/bin/bash +set -e + +# Transparently switch uid to root in order to perform the privileged function. +# SELinux restrictions and standard permissions still apply, the script and +# the calling user are restricted to being allowed to sudo only this. +if [ $(id -u) != 0 ]; then + exec sudo "$0" "$@" +fi + +ACTION="$1" + +case "$ACTION" in + start) + /bin/systemctl start ollama.service + ;; + stop) + /bin/systemctl stop ollama.service + ;; + *) + echo "Usage: $0 {start|stop}" >&2 + exit 2 + ;; +esac diff --git a/ic-os/components/guestos/ollama/ollama.service b/ic-os/components/guestos/ollama/ollama.service new file mode 100644 index 000000000000..721b1ea08830 --- /dev/null +++ b/ic-os/components/guestos/ollama/ollama.service @@ -0,0 +1,48 @@ +[Unit] +Description=Ollama LLM runtime (disabled by default) +Documentation=https://ollama.com +After=network-online.target +Wants=network-online.target + +# This service is intentionally disabled by default. It ships with the +# GuestOS image so that it can be started on demand with: +# systemctl enable --now ollama.service +# See the explicit `systemctl disable ollama.service` in the GuestOS +# Dockerfile which overrides the blanket enable loop. + +[Service] +Type=simple +User=ollama +Group=ollama + +# systemd will create /var/lib/ollama (mode 0750, owned by ollama:ollama) +# before ExecStart, and ensure it exists across upgrades. This avoids +# relying on the rootfs-baked directory, which is not visible on the +# separately-mounted /var filesystem. +StateDirectory=ollama +StateDirectoryMode=0750 + +# OLLAMA_MODELS points at the baked-in read-only model store on the +# dm-verity-signed root partition. This makes the pre-packaged models (e.g. +# gemma3:1b) available offline, but prevents `ollama pull` of new models at +# runtime. To allow pulling additional models, operators can override this +# with a drop-in (e.g. /etc/systemd/system/ollama.service.d/override.conf) +# pointing OLLAMA_MODELS at a writable location such as /var/lib/ollama/models +# and seeding it from /opt/ollama-models. +Environment=HOME=/var/lib/ollama +Environment=OLLAMA_MODELS=/opt/ollama-models +Environment=OLLAMA_HOST=0.0.0.0:11434 +Environment=LD_LIBRARY_PATH=/opt/ollama/lib/ollama + +ExecStart=/opt/ollama/bin/ollama serve + +Restart=on-failure +RestartSec=5s + +# No extra sandboxing: the GuestOS root is already dm-verity read-only and +# the orchestrator/replica services run without sandbox flags either. +# ProtectSystem/PrivateTmp/ReadWritePaths fail early in this environment +# with status=226/NAMESPACE because of the dm-crypt+LVM /var layout. + +[Install] +WantedBy=multi-user.target diff --git a/ic-os/guestos/context/BUILD.bazel b/ic-os/guestos/context/BUILD.bazel index 7c24c88e971e..d65a288bbe64 100644 --- a/ic-os/guestos/context/BUILD.bazel +++ b/ic-os/guestos/context/BUILD.bazel @@ -13,6 +13,10 @@ filegroup( "Dockerfile", "docker-base.dev", "docker-base.prod", + # Used by Dockerfile.base to pre-pull the gemma3:1b ollama model into + # /opt/ollama-models. Only consumed by the local-base-* envs which + # rebuild the base image from Dockerfile.base. + "ollama-pull-gemma.sh", "packages.common", "packages.dev", ], diff --git a/ic-os/guestos/context/Dockerfile b/ic-os/guestos/context/Dockerfile index b4a216633b00..cc348f208df5 100644 --- a/ic-os/guestos/context/Dockerfile +++ b/ic-os/guestos/context/Dockerfile @@ -113,7 +113,8 @@ RUN systemctl disable \ motd-news.service \ motd-news.timer \ fstrim.service \ - fstrim.timer + fstrim.timer \ + ollama.service # ------ GUESTOS WORK -------------------------------------------- @@ -126,6 +127,17 @@ RUN mkdir -p /var/lib/ic/backup \ /var/lib/ic/crypto \ /var/lib/ic/data +# Pre-create the directories where the ollama binary and CPU runtime +# libraries are injected at ext4-build time via the rootfs `extras` map in +# ic-os/guestos/defs.bzl. The published guestos-base image referenced by +# docker-base.{prod,dev} pre-dates the introduction of /opt/ollama, so +# without this `mkdir -p` the ext4 builder's `cp` step fails with +# "No such file or directory" for /opt/ollama/bin/ollama. Once the +# updated base image (with `mkdir -p /opt/ollama/...` in Dockerfile.base) +# is published and referenced from docker-base.{prod,dev}, this step +# becomes a no-op but is harmless. +RUN mkdir -p /opt/ollama/bin /opt/ollama/lib/ollama + # Create two mount points for temporary use during setup of "var" partition RUN mkdir -p /mnt/var_old /mnt/var_new @@ -231,6 +243,15 @@ RUN addgroup socks && \ adduser --system --disabled-password --shell /usr/sbin/nologin -c "Dante SOCKS Proxy" socks && \ adduser socks socks && chmod +s /usr/sbin/danted +# The "ollama" account. Used to run the `ollama serve` binary when the +# `ollama.service` unit is enabled (disabled by default; see +# /lib/systemd/system/ollama.service). +RUN addgroup --system ollama && \ + adduser --system --disabled-password --home /var/lib/ollama --shell /usr/sbin/nologin --ingroup ollama -c "Ollama" ollama && \ + mkdir -p /var/lib/ollama && \ + chown ollama:ollama /var/lib/ollama && \ + chmod 0750 /var/lib/ollama + # ------ INSTALL SCRIPTS ----------------------------------------- # Install IC binaries and other data late -- this means everything above diff --git a/ic-os/guestos/context/Dockerfile.base b/ic-os/guestos/context/Dockerfile.base index b2621f1f1430..c06967af3899 100644 --- a/ic-os/guestos/context/Dockerfile.base +++ b/ic-os/guestos/context/Dockerfile.base @@ -35,6 +35,19 @@ RUN cd /tmp/ && \ echo "c46e5b6f53948477ff3a19d97c58307394a29fe64a01905646f026ddc32cb65b node_exporter-1.10.2.linux-amd64.tar.gz" > node_exporter.sha256 && \ sha256sum -c node_exporter.sha256 +# Download and verify ollama. Ships the binary + shared libraries needed for +# CPU inference. Pinned to v0.21.1. The upstream bundle includes ~5 GiB of +# GPU runners (CUDA/MLX/Vulkan); these are stripped in the main Dockerfile +# because GuestOS nodes do CPU-only inference. +# To update: +# 1. Bump the version below. +# 2. Fetch the new sha256 with `curl -L https://github.com/ollama/ollama/releases/download/v/ollama-linux-amd64.tar.zst | sha256sum`. +# 3. Rebuild the base image and update docker-base.{prod,dev}. +RUN cd /tmp/ && \ + curl -L -O https://github.com/ollama/ollama/releases/download/v0.21.1/ollama-linux-amd64.tar.zst && \ + echo "528acc42750f755ce57b13b8138419ceac953b6fd488045f73dc1e80246da480 ollama-linux-amd64.tar.zst" > ollama.sha256 && \ + sha256sum -c ollama.sha256 + # # Second build stage: # - Download and cache minimal Ubuntu Server 24.04 LTS Docker image @@ -74,3 +87,34 @@ RUN cd /tmp/ && \ mkdir -p /etc/node_exporter && \ tar --strip-components=1 -C /usr/local/bin/ -zvxf node_exporter-1.10.2.linux-amd64.tar.gz node_exporter-1.10.2.linux-amd64/node_exporter && \ rm /tmp/node_exporter-1.10.2.linux-amd64.tar.gz + +# Install ollama (CPU-only). The upstream tarball is ~2 GiB and bundles GPU +# runners we don't use; we extract only the binary and the CPU ggml libraries +# into /opt/ollama. The resulting install is ~50 MiB. `curl` and `zstd` are +# already installed from packages.common above. +COPY --from=download /tmp/ollama-linux-amd64.tar.zst /tmp/ollama-linux-amd64.tar.zst +# The binary and CPU ggml .so files are *also* injected later via the ext4 +# extras path with explicit 0755 mode (see rootfs map in ic-os/guestos/defs.bzl). +# Podman's rootless squash/export drops the exec bit in this environment, so +# relying on Dockerfile.base-only install would yield a 0644 binary at runtime. +# We still need the binary here because `ollama-pull-gemma.sh` executes it to +# pre-populate /opt/ollama-models. +RUN mkdir -p /opt/ollama/bin /opt/ollama/lib/ollama /opt/ollama-models && \ + cd /tmp/ && \ + zstd -d ollama-linux-amd64.tar.zst -o ollama.tar && \ + tar -xf ollama.tar bin/ollama && \ + tar -tf ollama.tar | grep -E '^lib/ollama/libggml-(base|cpu-)[^/]*\.so' > /tmp/ollama-cpu-libs.txt && \ + tar -xf ollama.tar -T /tmp/ollama-cpu-libs.txt && \ + mv bin/ollama /opt/ollama/bin/ollama && \ + mv lib/ollama/libggml-*.so* /opt/ollama/lib/ollama/ && \ + chmod 0755 /opt/ollama/bin/ollama /opt/ollama/lib/ollama/*.so* && \ + rm -rf /tmp/bin /tmp/lib /tmp/ollama.tar /tmp/ollama-linux-amd64.tar.zst /tmp/ollama-cpu-libs.txt && \ + ln -s /opt/ollama/bin/ollama /usr/local/bin/ollama + +# Pre-pull the gemma3:1b model into /opt/ollama-models (read-only baked-in +# model store). The on-disk layout produced by `ollama pull` is +# /opt/ollama-models/{manifests,blobs}. Running +# `OLLAMA_MODELS=/opt/ollama-models ollama run gemma3:1b` at runtime will +# serve the model without any network access. +COPY ollama-pull-gemma.sh /tmp/ollama-pull-gemma.sh +RUN sh -eu /tmp/ollama-pull-gemma.sh && rm /tmp/ollama-pull-gemma.sh diff --git a/ic-os/guestos/context/ollama-pull-gemma.sh b/ic-os/guestos/context/ollama-pull-gemma.sh new file mode 100644 index 000000000000..6d005eedacd2 --- /dev/null +++ b/ic-os/guestos/context/ollama-pull-gemma.sh @@ -0,0 +1,36 @@ +#!/bin/sh +# Helper used from Dockerfile.base to pre-pull the Gemma model into the +# baked-in read-only model store at /opt/ollama-models. +# +# This is only invoked at base-image build time. See ../guestos/ollama.service +# for the runtime service definition. + +set -eu + +export OLLAMA_MODELS=/opt/ollama-models +export LD_LIBRARY_PATH=/opt/ollama/lib/ollama +export OLLAMA_HOST=127.0.0.1:11434 + +mkdir -p "${OLLAMA_MODELS}" + +/opt/ollama/bin/ollama serve >/tmp/ollama-serve.log 2>&1 & +OLLAMA_PID=$! +trap 'kill "${OLLAMA_PID}" 2>/dev/null || true; wait "${OLLAMA_PID}" 2>/dev/null || true' EXIT + +# Wait for the server to come up. +for i in $(seq 1 60); do + if curl -fs "http://${OLLAMA_HOST}/" >/dev/null 2>&1; then + break + fi + if [ "${i}" = 60 ]; then + echo "ollama server failed to start within 60s" >&2 + cat /tmp/ollama-serve.log >&2 || true + exit 1 + fi + sleep 1 +done + +/opt/ollama/bin/ollama pull gemma3:1b + +# Sanity-check that the model is usable offline. +/opt/ollama/bin/ollama list | grep -q '^gemma3:1b ' diff --git a/ic-os/guestos/defs.bzl b/ic-os/guestos/defs.bzl index 19ee4989fdcd..d512a313aafe 100644 --- a/ic-os/guestos/defs.bzl +++ b/ic-os/guestos/defs.bzl @@ -59,14 +59,34 @@ def image_deps(mode, malicious = False): # additional libraries to install "//rs/ic_os/release:nss_icos": "/usr/lib/x86_64-linux-gnu/libnss_icos.so.2:0644", # Allows referring to the guest IPv6 by name guestos from host, and host as hostos from guest. "//rs/ic_os/release:config_tool": "/opt/ic/bin/config_tool:0755", + + # Ollama CPU-only runtime. The binary and CPU ggml libraries are + # injected via the extras path (not via Dockerfile.base) so that + # they get explicit 0755 perms via `install -m` in the ext4 + # builder, bypassing rootless-podman's tendency to drop the + # exec bit during squash/export. The pre-pulled gemma3:1b model + # in /opt/ollama-models is still baked in via Dockerfile.base + # (model files are read-only 0644 data, no exec-bit concern). + "@ollama//:bin/ollama": "/opt/ollama/bin/ollama:0755", + "@ollama//:lib/ollama/libggml-base.so.0.0.0": "/opt/ollama/lib/ollama/libggml-base.so.0.0.0:0755", + "@ollama//:lib/ollama/libggml-cpu-alderlake.so": "/opt/ollama/lib/ollama/libggml-cpu-alderlake.so:0755", + "@ollama//:lib/ollama/libggml-cpu-haswell.so": "/opt/ollama/lib/ollama/libggml-cpu-haswell.so:0755", + "@ollama//:lib/ollama/libggml-cpu-icelake.so": "/opt/ollama/lib/ollama/libggml-cpu-icelake.so:0755", + "@ollama//:lib/ollama/libggml-cpu-sandybridge.so": "/opt/ollama/lib/ollama/libggml-cpu-sandybridge.so:0755", + "@ollama//:lib/ollama/libggml-cpu-skylakex.so": "/opt/ollama/lib/ollama/libggml-cpu-skylakex.so:0755", + "@ollama//:lib/ollama/libggml-cpu-sse42.so": "/opt/ollama/lib/ollama/libggml-cpu-sse42.so:0755", + "@ollama//:lib/ollama/libggml-cpu-x64.so": "/opt/ollama/lib/ollama/libggml-cpu-x64.so:0755", }, # Set various configuration values "container_context_files": Label("//ic-os/guestos/context:context-files"), "component_files": component_files(mode), "partition_table": Label("//ic-os/guestos:partitions.csv"), + # rootfs_size was bumped from 3G to 5G to accommodate the ollama + # binary + CPU runtime libraries (~50 MiB) and the pre-pulled + # gemma3:1b model (~815 MiB) baked into /opt/ollama-models. "expanded_size": "50G", - "rootfs_size": "3G", + "rootfs_size": "5G", "bootfs_size": "1G", "grub_config": Label("//ic-os/bootloader:guestos_grub.cfg"), diff --git a/ic-os/guestos/envs/local-base-dev/BUILD.bazel b/ic-os/guestos/envs/local-base-dev/BUILD.bazel index a74ed03e299e..3432a3d5824f 100644 --- a/ic-os/guestos/envs/local-base-dev/BUILD.bazel +++ b/ic-os/guestos/envs/local-base-dev/BUILD.bazel @@ -13,5 +13,8 @@ icos_build( "manual", "no-cache", ], - visibility = ["//rs:ic-os-pkg"], + visibility = [ + "//rs:ic-os-pkg", + "//rs:system-tests-pkg", + ], ) diff --git a/rs/ic_os/config/tool/templates/ic.json5.template b/rs/ic_os/config/tool/templates/ic.json5.template index 9432e3667299..eb89b4deecad 100644 --- a/rs/ic_os/config/tool/templates/ic.json5.template +++ b/rs/ic_os/config/tool/templates/ic.json5.template @@ -279,6 +279,8 @@ table ip6 filter {\n\ ip6 saddr { {{ ipv6_prefix }} } ct state { new } tcp dport { 7070, 9090, 9091, 9100, 19531, 19100, 19522 } accept\n\ # Allow access from HostOS metrics-proxy so GuestOS metrics-proxy can proxy certain metrics to HostOS\n\ ip6 saddr { hostos } ct state { new } tcp dport { 42372 } accept\n\ + # Allow access to the ollama HTTP API (disabled by default, started on-demand via systemctl enable --now ollama).\n\ + ip6 daddr { ::/0 } ct state { new } tcp dport { 11434 } accept\n\ # Custom templated rules\n\ <>\n\ <>\n\ @@ -418,6 +420,8 @@ table ip6 filter {\n\ ct state { established, related } accept\n\ ip6 saddr { {{ ipv6_prefix }} } ct state { new } tcp dport { 7070, 9091, 9100, 9324, 19531, 19100, 19522 } accept\n\ ip6 saddr { ::-ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff } ct state new tcp dport 443 accept\n\ + # Allow access to the ollama HTTP API (disabled by default, started on-demand via systemctl enable --now ollama).\n\ + ip6 daddr { ::/0 } ct state { new } tcp dport { 11434 } accept\n\ \n\ <>\n\ <>\n\ diff --git a/rs/orchestrator/src/ai_node.rs b/rs/orchestrator/src/ai_node.rs new file mode 100644 index 000000000000..52a24769abde --- /dev/null +++ b/rs/orchestrator/src/ai_node.rs @@ -0,0 +1,131 @@ +use crate::{ + error::{OrchestratorError, OrchestratorResult}, + registry_helper::RegistryHelper, +}; +use ic_logger::{ReplicaLogger, debug, info, warn}; +use ic_types::{NodeId, RegistryVersion}; +use std::{ + path::PathBuf, + sync::{Arc, RwLock}, +}; +use tokio::process::Command; + +/// Periodically checks the registry for an `AiNodeRecord` matching this node's +/// `NodeId` and reconciles the local `ollama.service` accordingly: +/// +/// * If a record exists for this node and ollama is not running, start it. +/// * If no record exists for this node and ollama is running, stop it. +/// +/// The unit is expected to be present (and disabled by default) on the image. +/// Starting/stopping is performed via `sudo /opt/ic/bin/manage-ollama.sh +/// {start,stop}` because the orchestrator runs as `ic-replica` and needs +/// root to invoke `systemctl`. +pub(crate) struct AiNodeManager { + registry: Arc, + node_id: NodeId, + logger: ReplicaLogger, + ic_binary_dir: PathBuf, + last_applied_version: Arc>, + /// Cached desired state. `None` until the first successful reconcile. + /// `Some(true)` means "ollama should be running" (record present). + desired_running: Option, +} + +impl AiNodeManager { + pub(crate) fn new( + registry: Arc, + node_id: NodeId, + ic_binary_dir: PathBuf, + logger: ReplicaLogger, + ) -> Self { + Self { + registry, + node_id, + logger, + ic_binary_dir, + last_applied_version: Default::default(), + desired_running: None, + } + } + + #[allow(dead_code)] + pub(crate) fn get_last_applied_version(&self) -> Arc> { + Arc::clone(&self.last_applied_version) + } + + /// Checks the registry and (un)starts `ollama.service` if needed. + pub(crate) async fn check_and_update(&mut self) { + let registry_version = self.registry.get_latest_version(); + debug!( + self.logger, + "Checking AiNodeRecord at registry version: {}", registry_version + ); + + let should_run = match self + .registry + .get_ai_node_record(self.node_id, registry_version) + { + Ok(Some(_)) => true, + Ok(None) => false, + Err(e) => { + warn!( + self.logger, + "Failed to read AiNodeRecord for {} at registry version {}: {}", + self.node_id, + registry_version, + e + ); + return; + } + }; + + // First-time reconcile: force an explicit transition regardless of cache, + // because on boot we do not know whether the unit was left running from a + // previous life. After that we only act on actual changes. + let need_transition = match self.desired_running { + None => true, + Some(prev) => prev != should_run, + }; + + if need_transition { + let action = if should_run { "start" } else { "stop" }; + match self.run_manage_ollama(action).await { + Ok(()) => { + info!( + self.logger, + "ollama.service {}ed (AiNodeRecord present: {})", action, should_run + ); + self.desired_running = Some(should_run); + } + Err(e) => { + warn!(self.logger, "Failed to {} ollama.service: {}", action, e); + return; + } + } + } + + *self.last_applied_version.write().unwrap() = registry_version; + } + + async fn run_manage_ollama(&self, action: &str) -> OrchestratorResult<()> { + let script = self.ic_binary_dir.join("manage-ollama.sh"); + let out = Command::new("sudo") + .arg(script.as_os_str()) + .arg(action) + .output() + .await + .map_err(|e| { + OrchestratorError::IoError(format!("failed to spawn manage-ollama.sh {action}"), e) + })?; + + if !out.status.success() { + return Err(OrchestratorError::UpgradeError(format!( + "manage-ollama.sh {action} failed: {:?} - stdout: {} - stderr: {}", + out.status, + String::from_utf8_lossy(&out.stdout).trim(), + String::from_utf8_lossy(&out.stderr).trim() + ))); + } + Ok(()) + } +} diff --git a/rs/orchestrator/src/lib.rs b/rs/orchestrator/src/lib.rs index 035fc0772b1d..e2ee596f202a 100644 --- a/rs/orchestrator/src/lib.rs +++ b/rs/orchestrator/src/lib.rs @@ -32,6 +32,7 @@ //! registry and writes them to disk for other components of the //! system to read. +mod ai_node; pub mod args; mod boundary_node; mod catch_up_package_provider; diff --git a/rs/orchestrator/src/orchestrator.rs b/rs/orchestrator/src/orchestrator.rs index 5d7824716203..612ddd224fff 100644 --- a/rs/orchestrator/src/orchestrator.rs +++ b/rs/orchestrator/src/orchestrator.rs @@ -1,4 +1,5 @@ use crate::{ + ai_node::AiNodeManager, args::OrchestratorArgs, boundary_node::BoundaryNodeManager, catch_up_package_provider::{CatchUpPackageProvider, LocalCUPReader}, @@ -80,6 +81,7 @@ pub struct Orchestrator { registration: Option, subnet_assignment: Arc>, ipv4_configurator: Option, + ai_node_manager: Option, task_tracker: TaskTracker, } @@ -347,7 +349,7 @@ impl Orchestrator { let ipv4_configurator = Ipv4Configurator::new( Arc::clone(®istry), Arc::clone(&metrics), - ic_binary_directory, + ic_binary_directory.clone(), logger.clone(), ); @@ -358,6 +360,13 @@ impl Orchestrator { logger.clone(), ); + let ai_node_manager = AiNodeManager::new( + Arc::clone(®istry), + node_id, + ic_binary_directory, + logger.clone(), + ); + let orchestrator_dashboard = Some(OrchestratorDashboard::new( Arc::clone(®istry), node_id, @@ -385,6 +394,7 @@ impl Orchestrator { registration: Some(registration), subnet_assignment, ipv4_configurator: Some(ipv4_configurator), + ai_node_manager: Some(ai_node_manager), task_tracker, }) } @@ -565,6 +575,7 @@ impl Orchestrator { mut ssh_access_manager: SshAccessManager, mut firewall: Firewall, mut ipv4_configurator: Ipv4Configurator, + mut ai_node_manager: AiNodeManager, cancellation_token: CancellationToken, ) { loop { @@ -588,6 +599,8 @@ impl Orchestrator { firewall.check_and_update(); // Check and update the network configuration ipv4_configurator.check_and_update().await; + // Check if this node should be running ollama based on AiNodeRecord + ai_node_manager.check_and_update().await; tokio::select! { _ = tokio::time::sleep(CHECK_INTERVAL_SECS) => {} _ = cancellation_token.cancelled() => break @@ -623,10 +636,11 @@ impl Orchestrator { ); } - if let (Some(ssh), Some(firewall), Some(ipv4_configurator)) = ( + if let (Some(ssh), Some(firewall), Some(ipv4_configurator), Some(ai_node_manager)) = ( self.ssh_access_manager.take(), self.firewall.take(), self.ipv4_configurator.take(), + self.ai_node_manager.take(), ) { self.task_tracker.spawn( "ssh_key_firewall_rules_ipv4_config", @@ -635,6 +649,7 @@ impl Orchestrator { ssh, firewall, ipv4_configurator, + ai_node_manager, cancellation_token.clone(), ), ); diff --git a/rs/orchestrator/src/registry_helper.rs b/rs/orchestrator/src/registry_helper.rs index 848c4fac4362..cd62796c4500 100644 --- a/rs/orchestrator/src/registry_helper.rs +++ b/rs/orchestrator/src/registry_helper.rs @@ -3,6 +3,7 @@ use ic_consensus_cup_utils::make_registry_cup; use ic_interfaces_registry::RegistryClient; use ic_logger::ReplicaLogger; use ic_protobuf::registry::{ + ai_node::v1::AiNodeRecord, api_boundary_node::v1::ApiBoundaryNodeRecord, firewall::v1::FirewallRuleSet, hostos_version::v1::HostosVersionRecord, @@ -11,6 +12,7 @@ use ic_protobuf::registry::{ subnet::v1::{SubnetRecord, SubnetType}, }; use ic_registry_client_helpers::{ + ai_node::AiNodeRegistry, api_boundary_node::ApiBoundaryNodeRegistry, firewall::FirewallRegistry, hostos_version::HostosRegistry, @@ -303,6 +305,17 @@ impl RegistryHelper { .map_err(OrchestratorError::RegistryClientError) } + /// Returns the `AiNodeRecord` for `node_id` at `version`, if any. + pub(crate) fn get_ai_node_record( + &self, + node_id: NodeId, + version: RegistryVersion, + ) -> OrchestratorResult> { + self.registry_client + .get_ai_node_record(node_id, version) + .map_err(OrchestratorError::RegistryClientError) + } + /// Return the DC ID where the current replica is located. pub fn dc_id(&self) -> Option { let registry_version = self.get_latest_version(); diff --git a/rs/orchestrator/testdata/nftables_assigned_cloud_engine.conf.golden b/rs/orchestrator/testdata/nftables_assigned_cloud_engine.conf.golden index 469935c12d58..18661d5b6470 100644 --- a/rs/orchestrator/testdata/nftables_assigned_cloud_engine.conf.golden +++ b/rs/orchestrator/testdata/nftables_assigned_cloud_engine.conf.golden @@ -97,6 +97,8 @@ table ip6 filter { ip6 saddr { ::/64 } ct state { new } tcp dport { 7070, 9090, 9091, 9100, 19531, 19100, 19522 } accept # Allow access from HostOS metrics-proxy so GuestOS metrics-proxy can proxy certain metrics to HostOS ip6 saddr { hostos } ct state { new } tcp dport { 42372 } accept + # Allow access to the ollama HTTP API (disabled by default, started on-demand via systemctl enable --now ollama). + ip6 daddr { ::/0 } ct state { new } tcp dport { 11434 } accept # Custom templated rules ip6 saddr {2001:db8:85a3::8a2e:1370:7334,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e5,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e6,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e7,a4c2:7f91:3db6:1e8c:5a4f:cc92:b37:6e41} ct state { new } tcp dport {22,2497,4100,8080} accept # Automatic node whitelisting ip6 saddr {::ffff:5.5.5.5} ct state { new } tcp dport {1005} accept # node_gwp4o-eaaaa-aaaaa-aaaap-2ai diff --git a/rs/orchestrator/testdata/nftables_assigned_replica.conf.golden b/rs/orchestrator/testdata/nftables_assigned_replica.conf.golden index c1ada72e8c3d..8d927c0fbd25 100644 --- a/rs/orchestrator/testdata/nftables_assigned_replica.conf.golden +++ b/rs/orchestrator/testdata/nftables_assigned_replica.conf.golden @@ -97,6 +97,8 @@ table ip6 filter { ip6 saddr { ::/64 } ct state { new } tcp dport { 7070, 9090, 9091, 9100, 19531, 19100, 19522 } accept # Allow access from HostOS metrics-proxy so GuestOS metrics-proxy can proxy certain metrics to HostOS ip6 saddr { hostos } ct state { new } tcp dport { 42372 } accept + # Allow access to the ollama HTTP API (disabled by default, started on-demand via systemctl enable --now ollama). + ip6 daddr { ::/0 } ct state { new } tcp dport { 11434 } accept # Custom templated rules ip6 saddr {3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e5,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e6,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e7,a4c2:7f91:3db6:1e8c:5a4f:cc92:b37:6e41} ct state { new } tcp dport {22,2497,4100,8080} accept # Automatic node whitelisting ip6 saddr {::ffff:5.5.5.5} ct state { new } tcp dport {1005} accept # node_gwp4o-eaaaa-aaaaa-aaaap-2ai diff --git a/rs/orchestrator/testdata/nftables_boundary_node_app_subnet.conf.golden b/rs/orchestrator/testdata/nftables_boundary_node_app_subnet.conf.golden index ee516706f637..95abe290f10c 100644 --- a/rs/orchestrator/testdata/nftables_boundary_node_app_subnet.conf.golden +++ b/rs/orchestrator/testdata/nftables_boundary_node_app_subnet.conf.golden @@ -89,6 +89,8 @@ table ip6 filter { ct state { established, related } accept ip6 saddr { ::/64 } ct state { new } tcp dport { 7070, 9091, 9100, 9324, 19531, 19100, 19522 } accept ip6 saddr { ::-ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff } ct state new tcp dport 443 accept + # Allow access to the ollama HTTP API (disabled by default, started on-demand via systemctl enable --now ollama). + ip6 daddr { ::/0 } ct state { new } tcp dport { 11434 } accept ip6 saddr {3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e5,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e6,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e7} ct state { new } tcp dport {1080} accept # nodes for SOCKS proxy ip6 saddr {::ffff:5.5.5.5} ct state { new } tcp dport {1005} accept # node_gol4s-2gsaq-aaaaa-aaaap-2ai diff --git a/rs/orchestrator/testdata/nftables_boundary_node_system_subnet.conf.golden b/rs/orchestrator/testdata/nftables_boundary_node_system_subnet.conf.golden index 7f2bb9660281..3a6e6d51ff2b 100644 --- a/rs/orchestrator/testdata/nftables_boundary_node_system_subnet.conf.golden +++ b/rs/orchestrator/testdata/nftables_boundary_node_system_subnet.conf.golden @@ -89,6 +89,8 @@ table ip6 filter { ct state { established, related } accept ip6 saddr { ::/64 } ct state { new } tcp dport { 7070, 9091, 9100, 9324, 19531, 19100, 19522 } accept ip6 saddr { ::-ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff } ct state new tcp dport 443 accept + # Allow access to the ollama HTTP API (disabled by default, started on-demand via systemctl enable --now ollama). + ip6 daddr { ::/0 } ct state { new } tcp dport { 11434 } accept ip6 saddr {a4c2:7f91:3db6:1e8c:5a4f:cc92:b37:6e41} ct state { new } tcp dport {1080} accept # nodes for SOCKS proxy ip6 saddr {::ffff:5.5.5.5} ct state { new } tcp dport {1005} accept # node_gwp4o-eaaaa-aaaaa-aaaap-2ai diff --git a/rs/orchestrator/testdata/nftables_unassigned_cloud_engine.conf.golden b/rs/orchestrator/testdata/nftables_unassigned_cloud_engine.conf.golden index 6c41426b2aae..f08bdb28e110 100644 --- a/rs/orchestrator/testdata/nftables_unassigned_cloud_engine.conf.golden +++ b/rs/orchestrator/testdata/nftables_unassigned_cloud_engine.conf.golden @@ -96,6 +96,8 @@ table ip6 filter { ip6 saddr { ::/64 } ct state { new } tcp dport { 7070, 9090, 9091, 9100, 19531, 19100, 19522 } accept # Allow access from HostOS metrics-proxy so GuestOS metrics-proxy can proxy certain metrics to HostOS ip6 saddr { hostos } ct state { new } tcp dport { 42372 } accept + # Allow access to the ollama HTTP API (disabled by default, started on-demand via systemctl enable --now ollama). + ip6 daddr { ::/0 } ct state { new } tcp dport { 11434 } accept # Custom templated rules ip6 saddr {2001:db8:85a3::8a2e:1370:7334,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e5,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e6,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e7,a4c2:7f91:3db6:1e8c:5a4f:cc92:b37:6e41} ct state { new } tcp dport {22,2497,4100,8080} accept # Automatic node whitelisting ip6 saddr {::ffff:5.5.5.5} ct state { new } tcp dport {1005} accept # node_gwp4o-eaaaa-aaaaa-aaaap-2ai diff --git a/rs/orchestrator/testdata/nftables_unassigned_replica.conf.golden b/rs/orchestrator/testdata/nftables_unassigned_replica.conf.golden index 259f21f415bc..2e78d5dcead5 100644 --- a/rs/orchestrator/testdata/nftables_unassigned_replica.conf.golden +++ b/rs/orchestrator/testdata/nftables_unassigned_replica.conf.golden @@ -96,6 +96,8 @@ table ip6 filter { ip6 saddr { ::/64 } ct state { new } tcp dport { 7070, 9090, 9091, 9100, 19531, 19100, 19522 } accept # Allow access from HostOS metrics-proxy so GuestOS metrics-proxy can proxy certain metrics to HostOS ip6 saddr { hostos } ct state { new } tcp dport { 42372 } accept + # Allow access to the ollama HTTP API (disabled by default, started on-demand via systemctl enable --now ollama). + ip6 daddr { ::/0 } ct state { new } tcp dport { 11434 } accept # Custom templated rules ip6 saddr {3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e5,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e6,3fda:92b7:4c1e:8a23:7d61:2f9c:ab42:19e7,a4c2:7f91:3db6:1e8c:5a4f:cc92:b37:6e41} ct state { new } tcp dport {22,2497,4100,8080} accept # Automatic node whitelisting ip6 saddr {::ffff:5.5.5.5} ct state { new } tcp dport {1005} accept # node_gwp4o-eaaaa-aaaaa-aaaap-2ai diff --git a/rs/protobuf/def/registry/ai_node/v1/ai_node.proto b/rs/protobuf/def/registry/ai_node/v1/ai_node.proto new file mode 100644 index 000000000000..5ae2e6773630 --- /dev/null +++ b/rs/protobuf/def/registry/ai_node/v1/ai_node.proto @@ -0,0 +1,13 @@ +syntax = "proto3"; + +package registry.ai_node.v1; + +// A record for a node that has been promoted to an "AI node". An AiNode is an +// extension of a regular `NodeRecord` (keyed by the same `NodeId`) and carries +// only AI-specific data. All operational data (IPs, domain, operator, etc.) is +// inherited from the underlying `NodeRecord`. +message AiNodeRecord { + // Optional principal id (bytes) of the subnet this AI node is associated + // with. When `None`, the node is not associated with any subnet. + optional bytes subnet_id = 1; +} diff --git a/rs/protobuf/generator/src/lib.rs b/rs/protobuf/generator/src/lib.rs index 94511a6bedbb..098fe201848b 100644 --- a/rs/protobuf/generator/src/lib.rs +++ b/rs/protobuf/generator/src/lib.rs @@ -265,9 +265,14 @@ fn build_registry_proto(def: &Path, out: &Path) { ".registry.api_boundary_node.v1.ApiBoundaryNodeRecord", "#[derive(serde::Serialize, serde::Deserialize)]", ); + config.type_attribute( + ".registry.ai_node.v1.AiNodeRecord", + "#[derive(serde::Serialize, serde::Deserialize)]", + ); let registry_files = [ def.join("registry/api_boundary_node/v1/api_boundary_node.proto"), + def.join("registry/ai_node/v1/ai_node.proto"), def.join("registry/crypto/v1/crypto.proto"), def.join("registry/node_operator/v1/node_operator.proto"), def.join("registry/nns/v1/nns.proto"), diff --git a/rs/protobuf/src/gen/registry/registry.ai_node.v1.rs b/rs/protobuf/src/gen/registry/registry.ai_node.v1.rs new file mode 100644 index 000000000000..6cf488151d79 --- /dev/null +++ b/rs/protobuf/src/gen/registry/registry.ai_node.v1.rs @@ -0,0 +1,12 @@ +// This file is @generated by prost-build. +/// A record for a node that has been promoted to an "AI node". An AiNode is an +/// extension of a regular `NodeRecord` (keyed by the same `NodeId`) and carries +/// only AI-specific data. All operational data (IPs, domain, operator, etc.) is +/// inherited from the underlying `NodeRecord`. +#[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message)] +pub struct AiNodeRecord { + /// Optional principal id (bytes) of the subnet this AI node is associated + /// with. When `None`, the node is not associated with any subnet. + #[prost(bytes = "vec", optional, tag = "1")] + pub subnet_id: ::core::option::Option<::prost::alloc::vec::Vec>, +} diff --git a/rs/protobuf/src/registry.rs b/rs/protobuf/src/registry.rs index fbbcfe8ffcf6..07a700e6cd4f 100644 --- a/rs/protobuf/src/registry.rs +++ b/rs/protobuf/src/registry.rs @@ -1,3 +1,4 @@ +pub mod ai_node; pub mod api_boundary_node; pub mod crypto; pub mod dc; diff --git a/rs/protobuf/src/registry/ai_node.rs b/rs/protobuf/src/registry/ai_node.rs new file mode 100644 index 000000000000..5e70a10285d5 --- /dev/null +++ b/rs/protobuf/src/registry/ai_node.rs @@ -0,0 +1,3 @@ +#[allow(clippy::all)] +#[path = "../gen/registry/registry.ai_node.v1.rs"] +pub mod v1; diff --git a/rs/registry/admin/bin/main.rs b/rs/registry/admin/bin/main.rs index 39a0ea9a9f35..f1ef0dfa2aa0 100644 --- a/rs/registry/admin/bin/main.rs +++ b/rs/registry/admin/bin/main.rs @@ -73,6 +73,7 @@ use ic_nns_init::make_hsm_sender; use ic_nns_test_utils::governance::{HardResetNnsRootToVersionPayload, UpgradeRootProposal}; use ic_protobuf::registry::replica_version::v1::GuestLaunchMeasurements; use ic_protobuf::registry::{ + ai_node::v1::AiNodeRecord, api_boundary_node::v1::ApiBoundaryNodeRecord, crypto::v1::{PublicKey, X509PublicKeyCert}, dc::v1::{AddOrRemoveDataCentersProposalPayload, DataCenterRecord}, @@ -92,10 +93,11 @@ use ic_registry_client_helpers::{ ecdsa_keys::EcdsaKeysRegistry, hostos_version::HostosRegistry, subnet::SubnetRegistry, }; use ic_registry_keys::{ - API_BOUNDARY_NODE_RECORD_KEY_PREFIX, FirewallRulesScope, NODE_OPERATOR_RECORD_KEY_PREFIX, - NODE_RECORD_KEY_PREFIX, NODE_REWARDS_TABLE_KEY, ROOT_SUBNET_ID_KEY, - get_node_operator_id_from_record_key, get_node_record_node_id, is_node_operator_record_key, - is_node_record_key, make_api_boundary_node_record_key, make_blessed_replica_versions_key, + AI_NODE_RECORD_KEY_PREFIX, API_BOUNDARY_NODE_RECORD_KEY_PREFIX, FirewallRulesScope, + NODE_OPERATOR_RECORD_KEY_PREFIX, NODE_RECORD_KEY_PREFIX, NODE_REWARDS_TABLE_KEY, + ROOT_SUBNET_ID_KEY, get_node_operator_id_from_record_key, get_node_record_node_id, + is_node_operator_record_key, is_node_record_key, make_ai_node_record_key, + make_api_boundary_node_record_key, make_blessed_replica_versions_key, make_canister_migrations_record_key, make_crypto_node_key, make_crypto_threshold_signing_pubkey_key, make_crypto_tls_cert_key, make_data_center_record_key, make_firewall_config_record_key, make_firewall_rules_record_key, @@ -126,6 +128,7 @@ use prost::Message; use recover_subnet::ProposeToUpdateRecoveryCupCmd; use registry_canister::mutations::{ complete_canister_migration::CompleteCanisterMigrationPayload, + do_add_ai_nodes::AddAiNodesPayload, do_add_api_boundary_nodes::AddApiBoundaryNodesPayload, do_add_node_operator::AddNodeOperatorPayload, do_change_subnet_membership::ChangeSubnetMembershipPayload, @@ -133,6 +136,7 @@ use registry_canister::mutations::{ do_deploy_guestos_to_all_subnet_nodes::DeployGuestosToAllSubnetNodesPayload, do_deploy_guestos_to_all_unassigned_nodes::DeployGuestosToAllUnassignedNodesPayload, do_migrate_node_operator_directly::MigrateNodeOperatorPayload, + do_remove_ai_nodes::RemoveAiNodesPayload, do_remove_api_boundary_nodes::RemoveApiBoundaryNodesPayload, do_remove_node_operators::RemoveNodeOperatorsPayload, do_revise_elected_replica_versions::ReviseElectedGuestosVersionsPayload, @@ -141,6 +145,7 @@ use registry_canister::mutations::{ NodeSshAccess, SetSubnetOperationalLevelPayload, operational_level, }, do_swap_node_in_subnet_directly::SwapNodeInSubnetDirectlyPayload, + do_update_ai_node_subnet::UpdateAiNodeSubnetPayload, do_update_api_boundary_nodes_version::DeployGuestosToSomeApiBoundaryNodes, do_update_elected_hostos_versions::ReviseElectedHostosVersionsPayload, do_update_node_operator_config::UpdateNodeOperatorConfigPayload, @@ -573,6 +578,22 @@ enum SubCommand { /// Migrate node operator for nodes directly, without governance. MigrateNodeOperatorDirectly(MigrateNodeOperatorDirectlyCmd), + /// Register AI nodes directly, without governance. + AddAiNodes(AddAiNodesCmd), + + /// Remove AI nodes directly, without governance. + RemoveAiNodes(RemoveAiNodesCmd), + + /// Update the subnet association of an AI node directly, without + /// governance. + UpdateAiNodeSubnet(UpdateAiNodeSubnetCmd), + + /// Retrieve an AI node record. + GetAiNode(GetAiNodeCmd), + + /// Retrieve all AI node ids. + GetAiNodes, + /// Update local registry store by pulling from remote URL UpdateRegistryLocalStore(UpdateRegistryLocalStoreCmd), @@ -3389,6 +3410,45 @@ struct MigrateNodeOperatorDirectlyCmd { pub new_node_operator_id: PrincipalId, } +#[derive(Parser)] +struct AddAiNodesCmd { + /// The nodes to register as AI nodes. + #[clap(long, required = true, num_args(1..), alias = "node-ids")] + pub nodes: Vec, + + /// Subnet to associate the AI nodes with. Required: every AI node must be + /// associated with a subnet at creation time. Use `update-ai-node-subnet` + /// later to change or clear the association. + #[clap(long, required = true)] + pub subnet_id: PrincipalId, +} + +#[derive(Parser)] +struct RemoveAiNodesCmd { + /// The AI nodes to remove. + #[clap(long, required = true, num_args(1..), alias = "node-ids")] + pub nodes: Vec, +} + +#[derive(Parser)] +struct UpdateAiNodeSubnetCmd { + /// The AI node whose subnet association should be updated. + #[clap(long)] + pub node_id: PrincipalId, + + /// New subnet association for the AI node. Omit to disassociate the + /// node from any subnet. + #[clap(long)] + pub subnet_id: Option, +} + +/// Sub-command to fetch an AI node record from the registry. +#[derive(Parser)] +struct GetAiNodeCmd { + /// The node id. + node_id: PrincipalId, +} + /// Sub-command to vote on a root proposal to upgrade the governance canister. #[derive(Parser)] struct VoteOnRootProposalToUpgradeGovernanceCanisterCmd { @@ -4597,6 +4657,11 @@ async fn main() { SubCommand::SubmitRootProposalToUpgradeGovernanceCanister(_) => (), SubCommand::SwapNodeInSubnetDirectly(_) => (), SubCommand::MigrateNodeOperatorDirectly(_) => (), + SubCommand::AddAiNodes(_) => (), + SubCommand::RemoveAiNodes(_) => (), + SubCommand::UpdateAiNodeSubnet(_) => (), + SubCommand::GetAiNode(_) => (), + SubCommand::GetAiNodes => (), SubCommand::VoteOnRootProposalToUpgradeGovernanceCanister(_) => (), _ => panic!( "Specifying a secret key or HSM is only supported for \ @@ -5405,6 +5470,33 @@ async fn main() { SubCommand::MigrateNodeOperatorDirectly(cmd) => { migrate_node_operator_directly(registry_canister, cmd).await; } + SubCommand::AddAiNodes(cmd) => { + add_ai_nodes(registry_canister, cmd).await; + } + SubCommand::RemoveAiNodes(cmd) => { + remove_ai_nodes(registry_canister, cmd).await; + } + SubCommand::UpdateAiNodeSubnet(cmd) => { + update_ai_node_subnet(registry_canister, cmd).await; + } + SubCommand::GetAiNode(cmd) => { + print_and_get_last_value::( + make_ai_node_record_key(cmd.node_id.into()) + .as_bytes() + .to_vec(), + ®istry_canister, + opts.json, + ) + .await; + } + SubCommand::GetAiNodes => { + let records = get_ai_node_ids(reachable_nns_urls.clone()); + println!( + "{}", + serde_json::to_string_pretty(&records) + .expect("Failed to serialize the records to JSON") + ); + } SubCommand::GetPendingRootProposalsToUpgradeGovernanceCanister => { get_pending_root_proposals_to_upgrade_governance_canister(make_canister_client( reachable_nns_urls, @@ -6655,6 +6747,34 @@ fn get_api_boundary_node_ids(nns_url: Vec) -> Vec { .collect::>() } +fn get_ai_node_ids(nns_url: Vec) -> Vec { + let registry_client = RegistryClientImpl::new( + Arc::new(NnsDataProvider::new( + tokio::runtime::Handle::current(), + nns_url, + )), + None, + ); + // maximum number of retries, let the user ctrl+c if necessary + registry_client + .try_polling_latest_version(usize::MAX) + .unwrap(); + let keys = registry_client + .get_key_family( + AI_NODE_RECORD_KEY_PREFIX, + registry_client.get_latest_version(), + ) + .unwrap(); + + keys.iter() + .map(|k| { + k.strip_prefix(AI_NODE_RECORD_KEY_PREFIX) + .unwrap() + .to_string() + }) + .collect::>() +} + fn print_routing_table(routing_table: &Vec<(CanisterIdRange, SubnetId)>, version: RegistryVersion) { println!("Routing table. Most recent version is {version}"); @@ -6962,6 +7082,86 @@ async fn migrate_node_operator_directly( } } +async fn add_ai_nodes(registry_canister: RegistryCanister, cmd: AddAiNodesCmd) { + let nonce = generate_nonce(); + let request = AddAiNodesPayload { + node_ids: cmd.nodes.into_iter().map(NodeId::from).collect(), + subnet_id: Some(SubnetId::from(cmd.subnet_id)), + }; + + let payload = Encode!(&request).expect("Failed to serialize add_ai_nodes request."); + + let agent = registry_canister.choose_random_agent(); + + match agent + .execute_update( + ®ISTRY_CANISTER_ID, + ®ISTRY_CANISTER_ID, + "add_ai_nodes", + payload, + nonce, + ) + .await + .unwrap() + { + Some(_) => println!("AI nodes added successfully."), + None => panic!("No response was received from add_ai_nodes"), + } +} + +async fn remove_ai_nodes(registry_canister: RegistryCanister, cmd: RemoveAiNodesCmd) { + let nonce = generate_nonce(); + let request = RemoveAiNodesPayload { + node_ids: cmd.nodes.into_iter().map(NodeId::from).collect(), + }; + + let payload = Encode!(&request).expect("Failed to serialize remove_ai_nodes request."); + + let agent = registry_canister.choose_random_agent(); + + match agent + .execute_update( + ®ISTRY_CANISTER_ID, + ®ISTRY_CANISTER_ID, + "remove_ai_nodes", + payload, + nonce, + ) + .await + .unwrap() + { + Some(_) => println!("AI nodes removed successfully."), + None => panic!("No response was received from remove_ai_nodes"), + } +} + +async fn update_ai_node_subnet(registry_canister: RegistryCanister, cmd: UpdateAiNodeSubnetCmd) { + let nonce = generate_nonce(); + let request = UpdateAiNodeSubnetPayload { + node_id: NodeId::from(cmd.node_id), + subnet_id: cmd.subnet_id.map(SubnetId::from), + }; + + let payload = Encode!(&request).expect("Failed to serialize update_ai_node_subnet request."); + + let agent = registry_canister.choose_random_agent(); + + match agent + .execute_update( + ®ISTRY_CANISTER_ID, + ®ISTRY_CANISTER_ID, + "update_ai_node_subnet", + payload, + nonce, + ) + .await + .unwrap() + { + Some(_) => println!("AI node subnet updated successfully."), + None => panic!("No response was received from update_ai_node_subnet"), + } +} + /// A client view of an NNS canister. struct NnsCanisterClient { /// The agent to talk to the IC. diff --git a/rs/registry/canister/canister/canister.rs b/rs/registry/canister/canister/canister.rs index 2483aa72935b..8a70942abafc 100644 --- a/rs/registry/canister/canister/canister.rs +++ b/rs/registry/canister/canister/canister.rs @@ -39,6 +39,7 @@ use registry_canister::{ init::RegistryCanisterInitPayload, mutations::{ complete_canister_migration::CompleteCanisterMigrationPayload, + do_add_ai_nodes::AddAiNodesPayload, do_add_api_boundary_nodes::AddApiBoundaryNodesPayload, do_add_node_operator::AddNodeOperatorPayload, do_add_nodes_to_subnet::AddNodesToSubnetPayload, @@ -49,6 +50,7 @@ use registry_canister::{ do_deploy_guestos_to_all_unassigned_nodes::DeployGuestosToAllUnassignedNodesPayload, do_migrate_node_operator_directly::MigrateNodeOperatorPayload, do_recover_subnet::RecoverSubnetPayload, + do_remove_ai_nodes::RemoveAiNodesPayload, do_remove_api_boundary_nodes::RemoveApiBoundaryNodesPayload, do_remove_node_operators::RemoveNodeOperatorsPayload, do_remove_nodes_from_subnet::RemoveNodesFromSubnetPayload, @@ -57,6 +59,7 @@ use registry_canister::{ do_set_subnet_operational_level::SetSubnetOperationalLevelPayload, do_split_subnet::SplitSubnetPayload, do_swap_node_in_subnet_directly::SwapNodeInSubnetDirectlyPayload, + do_update_ai_node_subnet::UpdateAiNodeSubnetPayload, do_update_api_boundary_nodes_version::{ DeployGuestosToSomeApiBoundaryNodes, UpdateApiBoundaryNodesVersionPayload, }, @@ -765,6 +768,76 @@ fn deploy_guestos_to_some_api_boundary_nodes_(payload: DeployGuestosToSomeApiBou recertify_registry(); } +#[unsafe(export_name = "canister_update add_ai_nodes")] +fn add_ai_nodes() { + // This method can be called by anyone. + println!( + "{}call: add_ai_nodes from: {}", + LOG_PREFIX, + dfn_core::api::caller() + ); + over(candid_one, |payload: AddAiNodesPayload| { + add_ai_nodes_(payload) + }); +} + +#[candid_method(update, rename = "add_ai_nodes")] +fn add_ai_nodes_(payload: AddAiNodesPayload) { + registry_mut().do_add_ai_nodes(payload); + recertify_registry(); +} + +#[unsafe(export_name = "canister_update remove_ai_nodes")] +fn remove_ai_nodes() { + // This method can be called by anyone. + println!( + "{}call: remove_ai_nodes from: {}", + LOG_PREFIX, + dfn_core::api::caller() + ); + over(candid_one, |payload: RemoveAiNodesPayload| { + remove_ai_nodes_(payload) + }); +} + +#[candid_method(update, rename = "remove_ai_nodes")] +fn remove_ai_nodes_(payload: RemoveAiNodesPayload) { + registry_mut().do_remove_ai_nodes(payload); + recertify_registry(); +} + +#[unsafe(export_name = "canister_update update_ai_node_subnet")] +fn update_ai_node_subnet() { + // This method can be called by anyone. + println!( + "{}call: update_ai_node_subnet from: {}", + LOG_PREFIX, + dfn_core::api::caller() + ); + over(candid_one, |payload: UpdateAiNodeSubnetPayload| { + update_ai_node_subnet_(payload) + }); +} + +#[candid_method(update, rename = "update_ai_node_subnet")] +fn update_ai_node_subnet_(payload: UpdateAiNodeSubnetPayload) { + registry_mut().do_update_ai_node_subnet(payload); + recertify_registry(); +} + +#[unsafe(export_name = "canister_query get_ai_node_ids")] +fn get_ai_node_ids() { + over(candid_one, |()| -> Result, String> { + get_ai_node_ids_() + }) +} + +#[candid_method(query, rename = "get_ai_node_ids")] +fn get_ai_node_ids_() -> Result, String> { + let ids = registry().get_ai_node_ids()?; + Ok(ids.into_iter().map(|n| n.get()).collect()) +} + #[unsafe(export_name = "canister_update remove_nodes")] fn remove_nodes() { check_caller_is_governance_and_log("remove_nodes"); diff --git a/rs/registry/canister/canister/registry.did b/rs/registry/canister/canister/registry.did index 876ffbbd11ef..4fa322053afe 100644 --- a/rs/registry/canister/canister/registry.did +++ b/rs/registry/canister/canister/registry.did @@ -8,6 +8,11 @@ // test_implementated_interface_matches_declared_interface_exactly (defined in // ./tests.rs) ensures that the implementation stays in sync with this file. +type AddAiNodesPayload = record { + node_ids : vec principal; + subnet_id : opt principal; +}; + type AddApiBoundaryNodesPayload = record { version : text; node_ids : vec principal; @@ -317,6 +322,8 @@ type RecoverSubnetPayload = record { time_ns : nat64; }; +type RemoveAiNodesPayload = record { node_ids : vec principal }; + type RemoveApiBoundaryNodesPayload = record { node_ids : vec principal }; type RemoveFirewallRulesPayload = record { @@ -375,6 +382,16 @@ type SubnetFeatures = record { type SubnetType = variant { application; verified_application; system; cloud_engine }; +type UpdateAiNodeSubnetPayload = record { + node_id : principal; + subnet_id : opt principal; +}; + +type GetAiNodeIdsResponse = variant { + Ok : vec principal; + Err : text; +}; + type UpdateApiBoundaryNodesVersionPayload = record { version : text; node_ids : vec principal; @@ -563,6 +580,7 @@ type UpdateSubnetAdminsPayload = record { }; service : { + add_ai_nodes : (AddAiNodesPayload) -> (); add_api_boundary_nodes : (AddApiBoundaryNodesPayload) -> (); add_firewall_rules : (AddFirewallRulesPayload) -> (); add_node : (AddNodePayload) -> (principal); @@ -582,6 +600,7 @@ service : { ) -> (); deploy_guestos_to_some_api_boundary_nodes : (DeployGuestosToSomeApiBoundaryNodes) -> (); deploy_hostos_to_some_nodes : (DeployHostosToSomeNodes) -> (); + get_ai_node_ids : () -> (GetAiNodeIdsResponse) query; get_api_boundary_node_ids : (GetApiBoundaryNodeIdsRequest) -> (GetApiBoundaryNodeIdsResponse) query; get_build_metadata : () -> (text) query; get_chunk : (GetChunkRequest) -> (GetChunkResponse) query; @@ -592,6 +611,7 @@ service : { migrate_node_operator_directly : (MigrateNodeOperatorPayload) -> (); prepare_canister_migration : (PrepareCanisterMigrationPayload) -> (); recover_subnet : (RecoverSubnetPayload) -> (); + remove_ai_nodes : (RemoveAiNodesPayload) -> (); remove_api_boundary_nodes : (RemoveApiBoundaryNodesPayload) -> (); remove_firewall_rules : (RemoveFirewallRulesPayload) -> (); remove_node_directly : (RemoveNodeDirectlyPayload) -> (); @@ -603,6 +623,7 @@ service : { revise_elected_replica_versions : (ReviseElectedGuestosVersionsPayload) -> (); set_firewall_config : (SetFirewallConfigPayload) -> (); set_subnet_operational_level : (SetSubnetOperationalLevelPayload) -> (); + update_ai_node_subnet : (UpdateAiNodeSubnetPayload) -> (); update_api_boundary_nodes_version : (UpdateApiBoundaryNodesVersionPayload) -> (); update_elected_hostos_versions : (UpdateElectedHostosVersionsPayload) -> (); revise_elected_hostos_versions : (ReviseElectedHostosVersionsPayload) -> (); diff --git a/rs/registry/canister/src/invariants/ai_node.rs b/rs/registry/canister/src/invariants/ai_node.rs new file mode 100644 index 000000000000..4693680219c2 --- /dev/null +++ b/rs/registry/canister/src/invariants/ai_node.rs @@ -0,0 +1,84 @@ +use std::collections::BTreeMap; + +use ic_base_types::{NodeId, PrincipalId, SubnetId}; +use ic_protobuf::registry::ai_node::v1::AiNodeRecord; +use ic_registry_keys::{AI_NODE_RECORD_KEY_PREFIX, get_ai_node_record_node_id}; +use prost::Message; + +use super::common::{ + InvariantCheckError, RegistrySnapshot, get_node_record_from_snapshot, + get_subnet_ids_from_snapshot, +}; + +/// Returns all AI node records from the snapshot. +fn get_ai_node_records_from_snapshot( + snapshot: &RegistrySnapshot, +) -> Result, InvariantCheckError> { + let mut result = BTreeMap::::new(); + for (key, value) in snapshot.iter() { + let key_str = match String::from_utf8(key.clone()) { + Ok(s) => s, + Err(_) => continue, + }; + + if !key_str.starts_with(AI_NODE_RECORD_KEY_PREFIX) { + continue; + } + + if let Some(principal_id) = get_ai_node_record_node_id(&key_str) { + let record = + AiNodeRecord::decode(value.as_slice()).map_err(|err| InvariantCheckError { + msg: format!("Failed to decode AiNodeRecord for key={key_str}: {err}"), + source: None, + })?; + result.insert(NodeId::from(principal_id), record); + } + } + Ok(result) +} + +/// Checks AiNode invariants: +/// * Every AiNodeRecord has a corresponding NodeRecord +/// * If `subnet_id` is set, it must refer to a subnet that exists in the +/// registry +pub(crate) fn check_ai_node_invariants( + snapshot: &RegistrySnapshot, +) -> Result<(), InvariantCheckError> { + let ai_node_records = get_ai_node_records_from_snapshot(snapshot)?; + let subnet_ids: std::collections::HashSet = + get_subnet_ids_from_snapshot(snapshot).into_iter().collect(); + + for (ai_node_id, record) in ai_node_records { + // NodeRecord must exist. + let node_record = get_node_record_from_snapshot(ai_node_id, snapshot)?; + if node_record.is_none() { + return Err(InvariantCheckError { + msg: format!( + "AI Node with id={ai_node_id} doesn't have a corresponding NodeRecord" + ), + source: None, + }); + } + + // If subnet_id is set, it must parse and must refer to an existing + // subnet in the snapshot. + if let Some(raw) = record.subnet_id { + let principal_id = + PrincipalId::try_from(raw.as_slice()).map_err(|err| InvariantCheckError { + msg: format!("AI Node with id={ai_node_id} has a malformed subnet_id: {err}"), + source: None, + })?; + let subnet_id = SubnetId::from(principal_id); + if !subnet_ids.contains(&subnet_id) { + return Err(InvariantCheckError { + msg: format!( + "AI Node with id={ai_node_id} references subnet_id={subnet_id} that does not exist" + ), + source: None, + }); + } + } + } + + Ok(()) +} diff --git a/rs/registry/canister/src/invariants/checks.rs b/rs/registry/canister/src/invariants/checks.rs index 02d4258c1533..49f77668af98 100644 --- a/rs/registry/canister/src/invariants/checks.rs +++ b/rs/registry/canister/src/invariants/checks.rs @@ -1,6 +1,7 @@ use crate::{ common::LOG_PREFIX, invariants::{ + ai_node::check_ai_node_invariants, api_boundary_node::check_api_boundary_node_invariants, assignment::check_node_assignment_invariants, common::RegistrySnapshot, @@ -102,6 +103,9 @@ impl Registry { // API Boundary Node invariant result = result.and(check_api_boundary_node_invariants(&snapshot)); + // AI Node invariant + result = result.and(check_ai_node_invariants(&snapshot)); + // HostOS version invariants result = result.and(check_hostos_version_invariants(&snapshot)); diff --git a/rs/registry/canister/src/invariants/mod.rs b/rs/registry/canister/src/invariants/mod.rs index cde989a790fe..57cbc2f8fa62 100644 --- a/rs/registry/canister/src/invariants/mod.rs +++ b/rs/registry/canister/src/invariants/mod.rs @@ -1,3 +1,4 @@ +mod ai_node; mod api_boundary_node; mod assignment; mod checks; diff --git a/rs/registry/canister/src/mutations/ai_node.rs b/rs/registry/canister/src/mutations/ai_node.rs new file mode 100644 index 000000000000..d0ea63345998 --- /dev/null +++ b/rs/registry/canister/src/mutations/ai_node.rs @@ -0,0 +1,62 @@ +use std::str::FromStr; + +use crate::{ + common::{LOG_PREFIX, key_family::get_key_family_iter}, + registry::Registry, +}; + +use ic_base_types::NodeId; +use ic_protobuf::registry::ai_node::v1::AiNodeRecord; +use ic_registry_keys::{AI_NODE_RECORD_KEY_PREFIX, make_ai_node_record_key}; +use ic_registry_transport::pb::v1::RegistryValue; +use ic_types::PrincipalId; +use prost::Message; + +impl Registry { + /// Get the AiNode record + pub fn get_ai_node_record(&self, node_id: NodeId) -> Option { + let RegistryValue { + value: ai_node_record_vec, + version: _, + deletion_marker: _, + timestamp_nanoseconds: _, + } = self.get( + &make_ai_node_record_key(node_id).into_bytes(), + self.latest_version(), + )?; + + Some(AiNodeRecord::decode(ai_node_record_vec.as_slice()).unwrap()) + } + + /// Get the AiNode record or panic on error with a message. + pub fn get_ai_node_or_panic(&self, node_id: NodeId) -> AiNodeRecord { + self.get_ai_node_record(node_id).unwrap_or_else(|| { + panic!("{LOG_PREFIX}ai_node record for {node_id:} not found in the registry.") + }) + } + + /// Get all AI node IDs. + pub fn get_ai_node_ids(&self) -> Result, String> { + let mut err_ids = Vec::new(); + + let ids: Vec = get_key_family_iter::(self, AI_NODE_RECORD_KEY_PREFIX) + .filter_map(|(id_str, _)| match PrincipalId::from_str(&id_str) { + Ok(principal_id) => Some(NodeId::from(principal_id)), + Err(_) => { + err_ids.push(id_str); + None + } + }) + .collect(); + + if err_ids.is_empty() { + Ok(ids) + } else { + let err_msg = format!( + "The following AI node IDs couldn't be parsed from registry: [{}]", + err_ids.join(", ") + ); + Err(err_msg) + } + } +} diff --git a/rs/registry/canister/src/mutations/common.rs b/rs/registry/canister/src/mutations/common.rs index 10b8e78b08e7..078941c233c5 100644 --- a/rs/registry/canister/src/mutations/common.rs +++ b/rs/registry/canister/src/mutations/common.rs @@ -4,7 +4,7 @@ use ic_protobuf::registry::{ subnet::v1::SubnetListRecord, unassigned_nodes_config::v1::UnassignedNodesConfigRecord, }; use ic_registry_keys::{ - make_api_boundary_node_record_key, make_node_record_key, + make_ai_node_record_key, make_api_boundary_node_record_key, make_node_record_key, make_unassigned_nodes_config_record_key, }; use prost::Message; @@ -35,6 +35,19 @@ pub(crate) fn check_api_boundary_nodes_exist(registry: &Registry, node_ids: &[No }); } +pub(crate) fn check_ai_nodes_exist(registry: &Registry, node_ids: &[NodeId]) { + let version = registry.latest_version(); + + node_ids.iter().copied().for_each(|node_id| { + let key = make_ai_node_record_key(node_id); + + let record = registry.get(key.as_bytes(), version); + if record.is_none() { + panic!("record not found"); + } + }); +} + pub(crate) fn check_replica_version_is_blessed(registry: &Registry, replica_version_id: &str) { let blessed_version_ids = registry.get_blessed_replica_version_ids(); assert!( diff --git a/rs/registry/canister/src/mutations/do_add_ai_nodes.rs b/rs/registry/canister/src/mutations/do_add_ai_nodes.rs new file mode 100644 index 000000000000..a2de54fc6361 --- /dev/null +++ b/rs/registry/canister/src/mutations/do_add_ai_nodes.rs @@ -0,0 +1,99 @@ +use candid::{CandidType, Deserialize}; +#[cfg(target_arch = "wasm32")] +use dfn_core::println; +use ic_base_types::{NodeId, PrincipalId, SubnetId}; +use ic_protobuf::registry::ai_node::v1::AiNodeRecord; +use ic_registry_keys::{make_ai_node_record_key, make_subnet_record_key}; +use ic_registry_transport::insert; +use prost::Message; +use serde::Serialize; + +use crate::{common::LOG_PREFIX, registry::Registry}; + +#[derive(Clone, Eq, PartialEq, Debug, CandidType, Deserialize, Serialize)] +pub struct AddAiNodesPayload { + pub node_ids: Vec, + /// Optional subnet id to associate the AI nodes with. When `None`, the + /// nodes are registered without being associated to any subnet. + pub subnet_id: Option, +} + +impl Registry { + /// Adds an AiNodeRecord to the registry for each of the given node ids. + pub fn do_add_ai_nodes(&mut self, payload: AddAiNodesPayload) { + println!("{LOG_PREFIX}do_add_ai_nodes: {payload:?}"); + + // Ensure payload is valid + self.validate_add_ai_nodes_payload(&payload); + + // Serialize the subnet_id (if any) to its principal-bytes form once. + let subnet_id_bytes = payload.subnet_id.map(|s| s.get().to_vec()); + + // Mutations to insert AiNodeRecord + let mutations = payload.node_ids.into_iter().map(|node_id| { + let key = make_ai_node_record_key(node_id); + insert( + key, + AiNodeRecord { + subnet_id: subnet_id_bytes.clone(), + } + .encode_to_vec(), + ) + }); + + self.maybe_apply_mutation_internal(mutations.collect()); + } + + fn validate_add_ai_nodes_payload(&self, payload: &AddAiNodesPayload) { + // Ensure there are no duplicates + let unique_count = payload + .node_ids + .iter() + .collect::>() + .len(); + if unique_count != payload.node_ids.len() { + panic!("there are duplicate nodes") + } + + // If a subnet_id is provided, it must exist in the registry. + if let Some(subnet_id) = payload.subnet_id { + let key = make_subnet_record_key(subnet_id); + if self.get(key.as_bytes(), self.latest_version()).is_none() { + panic!("subnet {subnet_id} does not exist"); + } + } + + for node_id in payload.node_ids.iter() { + // Ensure node exists + self.get_node_or_panic(*node_id); + + // Ensure record does not exist (the node is not already an AI node) + let key = make_ai_node_record_key(*node_id); + + let record = self.get(key.as_bytes(), self.latest_version()); + + if record.is_some() { + panic!("record exists: {node_id}"); + } + + // Ensure node is not assigned to a subnet + self.get_subnet_list_record().subnets.iter().for_each(|id| { + let id = + SubnetId::from(PrincipalId::try_from(id).expect("failed to parse subnet id")); + + self.get_subnet_or_panic(id) + .membership + .iter() + .for_each(|id| { + let id = NodeId::from( + PrincipalId::try_from(id).expect("failed to parse principal id"), + ); + + if *node_id == id { + panic!("node assigned to subnet: {node_id}"); + } + }) + }); + } + } +} diff --git a/rs/registry/canister/src/mutations/do_remove_ai_nodes.rs b/rs/registry/canister/src/mutations/do_remove_ai_nodes.rs new file mode 100644 index 000000000000..972b92550b19 --- /dev/null +++ b/rs/registry/canister/src/mutations/do_remove_ai_nodes.rs @@ -0,0 +1,38 @@ +use candid::{CandidType, Deserialize}; +#[cfg(target_arch = "wasm32")] +use dfn_core::println; +use ic_base_types::NodeId; +use ic_registry_keys::make_ai_node_record_key; +use ic_registry_transport::delete; +use serde::Serialize; + +use crate::{common::LOG_PREFIX, registry::Registry}; + +use super::common::check_ai_nodes_exist; + +#[derive(Clone, Eq, PartialEq, Debug, CandidType, Deserialize, Serialize)] +pub struct RemoveAiNodesPayload { + pub node_ids: Vec, +} + +impl Registry { + /// Remove a set of AiNodeRecords from the registry + pub fn do_remove_ai_nodes(&mut self, payload: RemoveAiNodesPayload) { + println!("{LOG_PREFIX}do_remove_ai_nodes: {payload:?}"); + + // Ensure payload is valid + self.validate_remove_ai_nodes_payload(&payload); + + // Mutations to remove AiNodeRecords + let mutations = payload.node_ids.into_iter().map(|node_id| { + let key = make_ai_node_record_key(node_id); + delete(key) + }); + + self.maybe_apply_mutation_internal(mutations.collect()) + } + + fn validate_remove_ai_nodes_payload(&self, payload: &RemoveAiNodesPayload) { + check_ai_nodes_exist(self, &payload.node_ids); + } +} diff --git a/rs/registry/canister/src/mutations/do_update_ai_node_subnet.rs b/rs/registry/canister/src/mutations/do_update_ai_node_subnet.rs new file mode 100644 index 000000000000..16b0d5af33a6 --- /dev/null +++ b/rs/registry/canister/src/mutations/do_update_ai_node_subnet.rs @@ -0,0 +1,49 @@ +use candid::{CandidType, Deserialize}; +#[cfg(target_arch = "wasm32")] +use dfn_core::println; +use ic_base_types::{NodeId, SubnetId}; +use ic_registry_keys::{make_ai_node_record_key, make_subnet_record_key}; +use ic_registry_transport::update; +use prost::Message; +use serde::Serialize; + +use crate::{common::LOG_PREFIX, registry::Registry}; + +use super::common::check_ai_nodes_exist; + +#[derive(Clone, Eq, PartialEq, Debug, CandidType, Deserialize, Serialize)] +pub struct UpdateAiNodeSubnetPayload { + pub node_id: NodeId, + /// New subnet association for the AI node. When `None`, the AI node is + /// disassociated from any subnet. + pub subnet_id: Option, +} + +impl Registry { + /// Updates the `subnet_id` of an existing AiNodeRecord. + pub fn do_update_ai_node_subnet(&mut self, payload: UpdateAiNodeSubnetPayload) { + println!("{LOG_PREFIX}do_update_ai_node_subnet: {payload:?}"); + + // Ensure payload is valid + self.validate_update_ai_node_subnet_payload(&payload); + + let key = make_ai_node_record_key(payload.node_id); + let mut ai_node = self.get_ai_node_or_panic(payload.node_id); + ai_node.subnet_id = payload.subnet_id.map(|s| s.get().to_vec()); + + self.maybe_apply_mutation_internal(vec![update(key, ai_node.encode_to_vec())]); + } + + fn validate_update_ai_node_subnet_payload(&self, payload: &UpdateAiNodeSubnetPayload) { + // AiNodeRecord must exist. + check_ai_nodes_exist(self, &[payload.node_id]); + + // If a subnet_id is provided, it must exist in the registry. + if let Some(subnet_id) = payload.subnet_id { + let key = make_subnet_record_key(subnet_id); + if self.get(key.as_bytes(), self.latest_version()).is_none() { + panic!("subnet {subnet_id} does not exist"); + } + } + } +} diff --git a/rs/registry/canister/src/mutations/mod.rs b/rs/registry/canister/src/mutations/mod.rs index 5d83737b549c..4ae1a5a3c1b2 100644 --- a/rs/registry/canister/src/mutations/mod.rs +++ b/rs/registry/canister/src/mutations/mod.rs @@ -1,6 +1,8 @@ +pub mod ai_node; pub mod api_boundary_node; pub mod common; pub mod complete_canister_migration; +pub mod do_add_ai_nodes; pub mod do_add_api_boundary_nodes; pub mod do_add_node_operator; pub mod do_add_nodes_to_subnet; @@ -15,6 +17,7 @@ pub mod do_deploy_guestos_to_all_unassigned_nodes; pub mod do_migrate_canisters; pub mod do_migrate_node_operator_directly; pub mod do_recover_subnet; +pub mod do_remove_ai_nodes; pub mod do_remove_api_boundary_nodes; pub mod do_remove_node_operators; pub mod do_remove_nodes_from_subnet; @@ -24,6 +27,7 @@ pub mod do_set_firewall_config; pub mod do_set_subnet_operational_level; pub mod do_split_subnet; pub mod do_swap_node_in_subnet_directly; +pub mod do_update_ai_node_subnet; pub mod do_update_api_boundary_nodes_version; pub mod do_update_elected_hostos_versions; pub mod do_update_node_directly; diff --git a/rs/registry/helpers/src/ai_node.rs b/rs/registry/helpers/src/ai_node.rs new file mode 100644 index 000000000000..429f38eb7aff --- /dev/null +++ b/rs/registry/helpers/src/ai_node.rs @@ -0,0 +1,73 @@ +use crate::deserialize_registry_value; +use ic_base_types::{NodeId, RegistryVersion, SubnetId}; +use ic_interfaces_registry::{RegistryClient, RegistryClientResult}; +use ic_protobuf::registry::ai_node::v1::AiNodeRecord; +use ic_registry_keys::{ + AI_NODE_RECORD_KEY_PREFIX, get_ai_node_record_node_id, make_ai_node_record_key, +}; +use ic_types::{PrincipalId, registry::RegistryClientError}; + +pub trait AiNodeRegistry { + /// Returns all AI node ids registered at `version`. + fn get_ai_node_ids(&self, version: RegistryVersion) + -> Result, RegistryClientError>; + + /// Returns the `AiNodeRecord` for `node_id` at `version`, if any. + fn get_ai_node_record( + &self, + node_id: NodeId, + version: RegistryVersion, + ) -> RegistryClientResult; + + /// Returns all AI node ids associated with the given `subnet_id` at + /// `version`. + fn get_ai_nodes_for_subnet( + &self, + subnet_id: SubnetId, + version: RegistryVersion, + ) -> Result, RegistryClientError>; +} + +impl AiNodeRegistry for T { + fn get_ai_node_ids( + &self, + version: RegistryVersion, + ) -> Result, RegistryClientError> { + let ai_node_record_keys = self.get_key_family(AI_NODE_RECORD_KEY_PREFIX, version)?; + let res = ai_node_record_keys + .iter() + .filter_map(|s| get_ai_node_record_node_id(s)) + .map(NodeId::from) + .collect(); + Ok(res) + } + + fn get_ai_node_record( + &self, + node_id: NodeId, + version: RegistryVersion, + ) -> RegistryClientResult { + let bytes = self.get_value(&make_ai_node_record_key(node_id), version); + deserialize_registry_value::(bytes) + } + + fn get_ai_nodes_for_subnet( + &self, + subnet_id: SubnetId, + version: RegistryVersion, + ) -> Result, RegistryClientError> { + let target = subnet_id.get(); + let ids = self.get_ai_node_ids(version)?; + let mut result = Vec::new(); + for id in ids { + if let Some(record) = self.get_ai_node_record(id, version)? + && let Some(raw) = record.subnet_id + && let Ok(principal) = PrincipalId::try_from(raw.as_slice()) + && principal == target + { + result.push(id); + } + } + Ok(result) + } +} diff --git a/rs/registry/helpers/src/lib.rs b/rs/registry/helpers/src/lib.rs index 194bbab057e8..e1ab9ce59077 100644 --- a/rs/registry/helpers/src/lib.rs +++ b/rs/registry/helpers/src/lib.rs @@ -3,6 +3,7 @@ //! Traits specific to a particular component (crypto comes to mind) will move //! to the respective crate/component at some point in the future. +pub mod ai_node; pub mod api_boundary_node; pub mod blessed_replica_version; pub mod chain_keys; diff --git a/rs/registry/keys/src/lib.rs b/rs/registry/keys/src/lib.rs index 0113181ef87b..08b031e174d1 100644 --- a/rs/registry/keys/src/lib.rs +++ b/rs/registry/keys/src/lib.rs @@ -22,6 +22,7 @@ pub const ROOT_SUBNET_ID_KEY: &str = "nns_subnet_id"; pub const NODE_REWARDS_TABLE_KEY: &str = "node_rewards_table"; const UNASSIGNED_NODES_CONFIG_RECORD_KEY: &str = "unassigned_nodes_config"; +pub const AI_NODE_RECORD_KEY_PREFIX: &str = "ai_node_record_"; pub const API_BOUNDARY_NODE_RECORD_KEY_PREFIX: &str = "api_boundary_node_"; pub const NODE_RECORD_KEY_PREFIX: &str = "node_record_"; pub const NODE_OPERATOR_RECORD_KEY_PREFIX: &str = "node_operator_record_"; @@ -294,6 +295,26 @@ pub fn get_api_boundary_node_record_node_id(key: &str) -> Option { } } +/// Makes a key for an AiNodeRecord registry entry. +pub fn make_ai_node_record_key(node_id: NodeId) -> String { + format!("{}{}", AI_NODE_RECORD_KEY_PREFIX, node_id.get()) +} + +/// Checks whether a given key is an AI node record key. +pub fn is_ai_node_record_key(key: &str) -> bool { + key.starts_with(AI_NODE_RECORD_KEY_PREFIX) +} + +/// Returns the node_id associated with a given ai_node_record key if the key +/// is, in fact, a valid ai_node_record_key. +pub fn get_ai_node_record_node_id(key: &str) -> Option { + if let Some(key) = key.strip_prefix(AI_NODE_RECORD_KEY_PREFIX) { + PrincipalId::from_str(key).ok() + } else { + None + } +} + /// Makes a key for a threshold signature public key entry for a subnet. pub fn make_crypto_threshold_signing_pubkey_key(subnet_id: SubnetId) -> String { format!("{CRYPTO_THRESHOLD_SIGNING_KEY_PREFIX}{subnet_id}") diff --git a/rs/tests/configure_icos.bzl b/rs/tests/configure_icos.bzl index b2713fd7dca6..867935ef42d9 100644 --- a/rs/tests/configure_icos.bzl +++ b/rs/tests/configure_icos.bzl @@ -126,6 +126,12 @@ def configure_icos(guestos, guestos_update, hostos, hostos_update, setupos): if guestos_version == True: # HEAD version guestos_local(suffix, "dev") + elif guestos_version == "local_base_dev": + # Rebuilds Dockerfile.base locally instead of using the pinned + # ghcr.io/dfinity/guestos-base-dev image. Slow (~15-30 min), but + # necessary whenever Dockerfile.base itself is modified (e.g. to + # bundle new third-party tooling). + guestos_local(suffix, "local-base-dev") elif guestos_version == "malicious": guestos_local(suffix, "dev-malicious") elif guestos_version == "recovery_dev": diff --git a/rs/tests/testnets/BUILD.bazel b/rs/tests/testnets/BUILD.bazel index ab20367c254e..8ec709a606d9 100644 --- a/rs/tests/testnets/BUILD.bazel +++ b/rs/tests/testnets/BUILD.bazel @@ -78,6 +78,12 @@ system_test( system_test_nns( name = "small", enable_metrics = True, + # Build the GuestOS base image from Dockerfile.base locally so that + # modifications to Dockerfile.base (for example, the bundled `ollama` + # binary + baked-in Gemma model) actually end up in the image used by + # this testnet. Without this, the testnet would start from the pinned + # upstream base image, which has none of those extras. + guestos = "local_base_dev", tags = [ "dynamic_testnet", "manual", @@ -146,6 +152,12 @@ system_test( system_test_nns( name = "small_nns", enable_metrics = True, + # Build the GuestOS base image from Dockerfile.base locally so that + # modifications to Dockerfile.base (for example, the bundled `ollama` + # binary + baked-in Gemma model) actually end up in the image used by + # this testnet. Without this, the testnet would start from the pinned + # upstream base image, which has none of those extras. + guestos = "local_base_dev", tags = [ "dynamic_testnet", "manual", diff --git a/third_party/BUILD.ollama.bazel b/third_party/BUILD.ollama.bazel new file mode 100644 index 000000000000..0117e224a481 --- /dev/null +++ b/third_party/BUILD.ollama.bazel @@ -0,0 +1,18 @@ +# Exposes the CPU-only bits of the upstream ollama-linux-amd64 tarball as +# individual file targets. The upstream archive bundles GPU runners +# (CUDA/MLX/Vulkan) under lib/ollama//; those are intentionally +# *not* exported here since GuestOS nodes run inference on the CPU only. + +package(default_visibility = ["//visibility:public"]) + +exports_files([ + "bin/ollama", + "lib/ollama/libggml-base.so.0.0.0", + "lib/ollama/libggml-cpu-alderlake.so", + "lib/ollama/libggml-cpu-haswell.so", + "lib/ollama/libggml-cpu-icelake.so", + "lib/ollama/libggml-cpu-sandybridge.so", + "lib/ollama/libggml-cpu-skylakex.so", + "lib/ollama/libggml-cpu-sse42.so", + "lib/ollama/libggml-cpu-x64.so", +])