feat: support suffixes for memory amounts

CHANGELOG.md

@@ -17,6 +17,7 @@
 * feat: `icp canister logs` to display the current canister logs
   * use `--follow` to continuously poll for new logs. `--interval <n>` to poll every `n` seconds
 * feat: Support `k`, `m`, `b`, `t` suffixes in `.yaml` files when specifying cycles amounts
+* feat: Support `kb`, `kib`, `mb`, `mib`, `gb`, `gib` suffixes in `.yaml` files and CLI arguments when specifying memory amounts
 * feat: Add an optional root-key argument to canister commands
 * chore!: new passwords for identity encryption need to be at least 8 characters long
 * feat: Anonymous usage telemetry — collects command name, arguments, duration, and outcome

crates/icp-cli/src/commands/canister/create.rs

@@ -2,7 +2,7 @@ use anyhow::anyhow;
 use candid::{Nat, Principal};
 use clap::Args;
 use icp::context::Context;
-use icp::parsers::CyclesAmount;
+use icp::parsers::{CyclesAmount, MemoryAmount};
 use icp::{Canister, context::CanisterSelection, prelude::*};
 use icp_canister_interfaces::cycles_ledger::CanisterSettingsArg;
 
@@ -21,8 +21,9 @@ pub(crate) struct CanisterSettings {
     pub(crate) compute_allocation: Option<u64>,
 
     /// Optional memory allocation in bytes. If unset, memory is allocated dynamically.
+    /// Supports suffixes: kb, kib, mb, mib, gb, gib (e.g. "4gib" or "2.5kb").
     #[arg(long)]
-    pub(crate) memory_allocation: Option<u64>,
+    pub(crate) memory_allocation: Option<MemoryAmount>,
 
     /// Optional freezing threshold in seconds. Controls how long a canister can be inactive before being frozen.
     #[arg(long)]
@@ -86,8 +87,9 @@ impl CreateArgs {
             memory_allocation: self
                 .settings
                 .memory_allocation
-                .or(default.settings.memory_allocation)
-                .map(Nat::from),
+                .clone()
+                .or(default.settings.memory_allocation.clone())
+                .map(|m| Nat::from(m.get())),
             compute_allocation: self
                 .settings
                 .compute_allocation

crates/icp-cli/src/commands/canister/settings/update.rs

@@ -1,13 +1,12 @@
 use anyhow::bail;
-use byte_unit::{Byte, Unit};
 use clap::{ArgAction, Args};
 use dialoguer::Confirm;
 use ic_agent::Identity;
 use ic_agent::export::Principal;
 use ic_management_canister_types::{CanisterStatusResult, EnvironmentVariable, LogVisibility};
 use icp::ProjectLoadError;
 use icp::context::{CanisterSelection, Context, TermWriter};
-use icp::parsers::CyclesAmount;
+use icp::parsers::{CyclesAmount, MemoryAmount};
 use std::collections::{HashMap, HashSet};
 use std::io::Write;
 
@@ -97,8 +96,9 @@ pub(crate) struct UpdateArgs {
     #[arg(long, value_parser = compute_allocation_parser)]
     compute_allocation: Option<u8>,
 
-    #[arg(long, value_parser = memory_parser)]
-    memory_allocation: Option<Byte>,
+    /// Memory allocation in bytes. Supports suffixes: kb, kib, mb, mib, gb, gib (e.g. "4gib" or "2.5kb").
+    #[arg(long)]
+    memory_allocation: Option<MemoryAmount>,
 
     #[arg(long, value_parser = freezing_threshold_parser)]
     freezing_threshold: Option<u64>,
@@ -109,11 +109,13 @@ pub(crate) struct UpdateArgs {
     #[arg(long)]
     reserved_cycles_limit: Option<CyclesAmount>,
 
-    #[arg(long, value_parser = memory_parser)]
-    wasm_memory_limit: Option<Byte>,
+    /// Wasm memory limit in bytes. Supports suffixes: kb, kib, mb, mib, gb, gib (e.g. "4gib" or "2.5kb").
+    #[arg(long)]
+    wasm_memory_limit: Option<MemoryAmount>,
 
-    #[arg(long, value_parser = memory_parser)]
-    wasm_memory_threshold: Option<Byte>,
+    /// Wasm memory threshold in bytes. Supports suffixes: kb, kib, mb, mib, gb, gib (e.g. "4gib" or "2.5kb").
+    #[arg(long)]
+    wasm_memory_threshold: Option<MemoryAmount>,
 
     #[command(flatten)]
     log_visibility: Option<LogVisibilityOpt>,
@@ -221,13 +223,13 @@ pub(crate) async fn exec(ctx: &Context, args: &UpdateArgs) -> Result<(), anyhow:
         }
         update = update.with_compute_allocation(compute_allocation);
     }
-    if let Some(memory_allocation) = args.memory_allocation {
+    if let Some(memory_allocation) = &args.memory_allocation {
         if configured_settings.memory_allocation.is_some() {
             ctx.term.write_line(
                 "Warning: Memory allocation is already set in icp.yaml; this new value will be overridden on next settings sync"
             )?
         }
-        update = update.with_memory_allocation(memory_allocation.as_u64());
+        update = update.with_memory_allocation(memory_allocation.get());
     }
     if let Some(freezing_threshold) = args.freezing_threshold {
         if configured_settings.freezing_threshold.is_some() {
@@ -245,21 +247,21 @@ pub(crate) async fn exec(ctx: &Context, args: &UpdateArgs) -> Result<(), anyhow:
         }
         update = update.with_reserved_cycles_limit(reserved_cycles_limit.get());
     }
-    if let Some(wasm_memory_limit) = args.wasm_memory_limit {
+    if let Some(wasm_memory_limit) = &args.wasm_memory_limit {
         if configured_settings.wasm_memory_limit.is_some() {
             ctx.term.write_line(
                 "Warning: Wasm memory limit is already set in icp.yaml; this new value will be overridden on next settings sync"
             )?
         }
-        update = update.with_wasm_memory_limit(wasm_memory_limit.as_u64());
+        update = update.with_wasm_memory_limit(wasm_memory_limit.get());
     }
-    if let Some(wasm_memory_threshold) = args.wasm_memory_threshold {
+    if let Some(wasm_memory_threshold) = &args.wasm_memory_threshold {
         if configured_settings.wasm_memory_threshold.is_some() {
             ctx.term.write_line(
                 "Warning: Wasm memory threshold is already set in icp.yaml; this new value will be overridden on next settings sync"
             )?
         }
-        update = update.with_wasm_memory_threshold(wasm_memory_threshold.as_u64());
+        update = update.with_wasm_memory_threshold(wasm_memory_threshold.get());
     }
     if let Some(log_visibility) = log_visibility {
         if configured_settings.log_visibility.is_some() {
@@ -286,16 +288,6 @@ fn compute_allocation_parser(compute_allocation: &str) -> Result<u8, String> {
     Err("Must be a percent between 0 and 100".to_string())
 }
 
-fn memory_parser(memory_allocation: &str) -> Result<Byte, String> {
-    let limit = Byte::from_u64_with_unit(256, Unit::TiB).expect("256 TiB is a valid byte unit");
-    if let Ok(byte) = memory_allocation.parse::<Byte>()
-        && byte <= limit
-    {
-        return Ok(byte);
-    }
-    Err("Must be a value between 0..256 TiB inclusive, (e.g. '2GiB')".to_string())
-}
-
 fn freezing_threshold_parser(freezing_threshold: &str) -> Result<u64, String> {
     if let Ok(num) = freezing_threshold.parse::<u64>() {
         return Ok(num);

crates/icp-cli/src/operations/settings.rs

@@ -10,6 +10,7 @@ use ic_management_canister_types::{EnvironmentVariable, LogVisibility as IcLogVi
 use ic_utils::interfaces::ManagementCanister;
 use icp::{Canister, canister::Settings, context::TermWriter};
 use itertools::Itertools;
+use num_traits::ToPrimitive;
 use snafu::{ResultExt, Snafu};
 
 use crate::progress::{ProgressManager, ProgressManagerSettings};
@@ -79,11 +80,11 @@ pub(crate) async fn sync_settings(
     let &Settings {
         ref log_visibility,
         compute_allocation,
-        memory_allocation,
+        ref memory_allocation,
         freezing_threshold,
         ref reserved_cycles_limit,
-        wasm_memory_limit,
-        wasm_memory_threshold,
+        ref wasm_memory_limit,
+        ref wasm_memory_threshold,
         ref environment_variables,
     } = &canister.settings;
     let current_settings = status.settings;
@@ -114,13 +115,22 @@ pub(crate) async fn sync_settings(
         .as_ref()
         .is_none_or(|s| log_visibility_eq(s, &current_settings.log_visibility))
         && compute_allocation.is_none_or(|s| s == current_settings.compute_allocation)
-        && memory_allocation.is_none_or(|s| s == current_settings.memory_allocation)
+        && memory_allocation
+            .as_ref()
+            .map(|m| m.get())
+            .is_none_or(|s| current_settings.memory_allocation.0.to_u64() == Some(s))
         && freezing_threshold.is_none_or(|s| s == current_settings.freezing_threshold)
         && reserved_cycles_limit
             .as_ref()
             .is_none_or(|s| s.get() == current_settings.reserved_cycles_limit)
-        && wasm_memory_limit.is_none_or(|s| s == current_settings.wasm_memory_limit)
-        && wasm_memory_threshold.is_none_or(|s| s == current_settings.wasm_memory_threshold)
+        && wasm_memory_limit
+            .as_ref()
+            .map(|m| m.get())
+            .is_none_or(|s| current_settings.wasm_memory_limit.0.to_u64() == Some(s))
+        && wasm_memory_threshold
+            .as_ref()
+            .map(|m| m.get())
+            .is_none_or(|s| current_settings.wasm_memory_threshold.0.to_u64() == Some(s))
         && environment_variable_setting
             .as_ref()
             .is_none_or(|s| environment_variables_eq(s, &current_settings.environment_variables))
@@ -131,11 +141,11 @@ pub(crate) async fn sync_settings(
     mgmt.update_settings(cid)
         .with_optional_log_visibility(log_visibility_setting)
         .with_optional_compute_allocation(compute_allocation)
-        .with_optional_memory_allocation(memory_allocation)
+        .with_optional_memory_allocation(memory_allocation.as_ref().map(|m| m.get()))
         .with_optional_freezing_threshold(freezing_threshold)
         .with_optional_reserved_cycles_limit(reserved_cycles_limit.as_ref().map(|r| r.get()))
-        .with_optional_wasm_memory_limit(wasm_memory_limit)
-        .with_optional_wasm_memory_threshold(wasm_memory_threshold)
+        .with_optional_wasm_memory_limit(wasm_memory_limit.as_ref().map(|m| m.get()))
+        .with_optional_wasm_memory_threshold(wasm_memory_threshold.as_ref().map(|m| m.get()))
         .with_optional_environment_variables(environment_variable_setting)
         .build()
         .context(ValidateSettingsSnafu {

crates/icp-cli/tests/canister_create_tests.rs

@@ -185,15 +185,15 @@ async fn canister_create_with_settings() {
 }
 
 #[tokio::test]
-async fn canister_create_with_settings_reserved_cycles_suffix_in_yaml() {
+async fn canister_create_with_settings_suffix_in_yaml() {
     let ctx = TestContext::new();
 
     let project_dir = ctx.create_project_dir("icp");
 
     let f = NamedTempFile::new().expect("failed to create temporary file");
     let path = f.path();
 
-    // reserved_cycles_limit with suffix
+    // reserved_cycles_limit, memory_allocation and wasm_memory_limit with suffixes
     let pm = formatdoc! {r#"
             canisters:
               - name: my-canister
@@ -204,6 +204,8 @@ async fn canister_create_with_settings_reserved_cycles_suffix_in_yaml() {
                 settings:
                   compute_allocation: 1
                   reserved_cycles_limit: 1.2t
+                  memory_allocation: 2gib
+                  wasm_memory_limit: 0.25kib
 
             {NETWORK_RANDOM_PORT}
             {ENVIRONMENT_RANDOM_PORT}
@@ -245,8 +247,36 @@ async fn canister_create_with_settings_reserved_cycles_suffix_in_yaml() {
         .stdout(
             starts_with("Canister Id:")
                 .and(contains("Status: Running"))
-                .and(contains("Reserved cycles limit: 1_200_000_000_000")),
+                .and(contains("Reserved cycles limit: 1_200_000_000_000"))
+                .and(contains("Memory allocation: 2_147_483_648")),
         );
+
+    // Sync settings from manifest to apply wasm_memory_limit (not sent at create)
+    ctx.icp()
+        .current_dir(&project_dir)
+        .args([
+            "canister",
+            "settings",
+            "sync",
+            "my-canister",
+            "--environment",
+            "random-environment",
+        ])
+        .assert()
+        .success();
+
+    ctx.icp()
+        .current_dir(&project_dir)
+        .args([
+            "canister",
+            "status",
+            "my-canister",
+            "--environment",
+            "random-environment",
+        ])
+        .assert()
+        .success()
+        .stdout(contains("Wasm memory limit: 256"));
 }
 
 #[tokio::test]

crates/icp/src/canister/mod.rs

@@ -5,6 +5,8 @@ use icp_canister_interfaces::cycles_ledger::CanisterSettingsArg;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
+use crate::parsers::{CyclesAmount, MemoryAmount};
+
 pub mod build;
 pub mod recipe;
 pub mod sync;
@@ -208,7 +210,8 @@ pub struct Settings {
     pub compute_allocation: Option<u64>,
 
     /// Memory allocation in bytes. If unset, memory is allocated dynamically.
-    pub memory_allocation: Option<u64>,
+    /// Supports suffixes in YAML: kb, kib, mb, mib, gb, gib (e.g. "4gib" or "2.5kb").
+    pub memory_allocation: Option<MemoryAmount>,
 
     /// Freezing threshold in seconds. Controls how long a canister can be inactive before being frozen.
     pub freezing_threshold: Option<u64>,
@@ -217,13 +220,15 @@ pub struct Settings {
     /// Memory allocations that would push the reserved balance above this limit will fail.
     /// Supports suffixes in YAML: k, m, b, t (e.g. "4t" or "4.3t").
     #[serde(default)]
-    pub reserved_cycles_limit: Option<crate::parsers::CyclesAmount>,
+    pub reserved_cycles_limit: Option<CyclesAmount>,
 
     /// Wasm memory limit in bytes. Sets an upper bound for Wasm heap growth.
-    pub wasm_memory_limit: Option<u64>,
+    /// Supports suffixes in YAML: kb, kib, mb, mib, gb, gib (e.g. "4gib" or "2.5kb").
+    pub wasm_memory_limit: Option<MemoryAmount>,
 
     /// Wasm memory threshold in bytes. Triggers a callback when exceeded.
-    pub wasm_memory_threshold: Option<u64>,
+    /// Supports suffixes in YAML: kb, kib, mb, mib, gb, gib (e.g. "4gib" or "2.5kb").
+    pub wasm_memory_threshold: Option<MemoryAmount>,
 
     /// Environment variables for the canister as key-value pairs.
     /// These variables are accessible within the canister and can be used to configure
@@ -238,7 +243,7 @@ impl From<Settings> for CanisterSettingsArg {
             controllers: None,
             reserved_cycles_limit: settings.reserved_cycles_limit.map(|c| Nat::from(c.get())),
             log_visibility: settings.log_visibility.map(Into::into),
-            memory_allocation: settings.memory_allocation.map(Nat::from),
+            memory_allocation: settings.memory_allocation.map(|m| Nat::from(m.get())),
             compute_allocation: settings.compute_allocation.map(Nat::from),
         }
     }
@@ -356,6 +361,36 @@ allowed_viewers:
         );
     }
 
+    #[test]
+    fn settings_memory_allocation_parses_suffix() {
+        let yaml = "memory_allocation: 4gib";
+        let settings: Settings = serde_yaml::from_str(yaml).unwrap();
+        assert_eq!(
+            settings.memory_allocation.as_ref().map(|m| m.get()),
+            Some(4 * 1024 * 1024 * 1024)
+        );
+    }
+
+    #[test]
+    fn settings_memory_allocation_parses_number() {
+        let yaml = "memory_allocation: 4294967296";
+        let settings: Settings = serde_yaml::from_str(yaml).unwrap();
+        assert_eq!(
+            settings.memory_allocation.as_ref().map(|m| m.get()),
+            Some(4294967296)
+        );
+    }
+
+    #[test]
+    fn settings_wasm_memory_limit_parses_suffix() {
+        let yaml = "wasm_memory_limit: 1.5gib";
+        let settings: Settings = serde_yaml::from_str(yaml).unwrap();
+        assert_eq!(
+            settings.wasm_memory_limit.as_ref().map(|m| m.get()),
+            Some(1610612736)
+        );
+    }
+
     #[test]
     fn log_visibility_conversion_to_ic_type() {
         let controllers = LogVisibility::Controllers;

crates/icp/src/manifest/canister.rs

@@ -343,13 +343,16 @@ mod tests {
     use indoc::indoc;
     use std::collections::HashMap;
 
-    use crate::manifest::{
-        adapter::{
-            assets,
-            prebuilt::{self, RemoteSource, SourceField},
-            script,
+    use crate::{
+        manifest::{
+            adapter::{
+                assets,
+                prebuilt::{self, RemoteSource, SourceField},
+                script,
+            },
+            recipe::RecipeType,
         },
-        recipe::RecipeType,
+        parsers::MemoryAmount,
     };
 
     use super::*;
@@ -647,7 +650,7 @@ mod tests {
                 name: "my-canister".to_string(),
                 settings: Settings {
                     compute_allocation: Some(3),
-                    memory_allocation: Some(4294967296),
+                    memory_allocation: Some(MemoryAmount::from(4294967296)),
                     ..Default::default()
                 },
                 init_args: None,