[CLI] Add --memory option to apply and offer (#3461)

Author: un-def

src/dstack/_internal/cli/commands/offer.py

@@ -3,11 +3,11 @@
 from typing import List, Literal, cast
 
 from dstack._internal.cli.commands import APIBaseCommand
-from dstack._internal.cli.services.args import cpu_spec, disk_spec, gpu_spec
 from dstack._internal.cli.services.configurators.run import (
     BaseRunConfigurator,
 )
 from dstack._internal.cli.services.profile import register_profile_args
+from dstack._internal.cli.services.resources import register_resources_args
 from dstack._internal.cli.utils.common import console
 from dstack._internal.cli.utils.gpu import print_gpu_json, print_gpu_table
 from dstack._internal.cli.utils.run import print_offers_json, print_run_plan
@@ -47,29 +47,7 @@ def register_args(cls, parser: argparse.ArgumentParser):
             default=50,
         )
         cls.register_env_args(configuration_group)
-        configuration_group.add_argument(
-            "--cpu",
-            type=cpu_spec,
-            help="Request CPU for the run. "
-            "The format is [code]ARCH[/]:[code]COUNT[/] (all parts are optional)",
-            dest="cpu_spec",
-            metavar="SPEC",
-        )
-        configuration_group.add_argument(
-            "--gpu",
-            type=gpu_spec,
-            help="Request GPU for the run. "
-            "The format is [code]NAME[/]:[code]COUNT[/]:[code]MEMORY[/] (all parts are optional)",
-            dest="gpu_spec",
-            metavar="SPEC",
-        )
-        configuration_group.add_argument(
-            "--disk",
-            type=disk_spec,
-            help="Request the size range of disk for the run. Example [code]--disk 100GB..[/].",
-            metavar="RANGE",
-            dest="disk_spec",
-        )
+        register_resources_args(configuration_group)
         register_profile_args(parser)
 
 

src/dstack/_internal/cli/services/configurators/run.py

@@ -12,8 +12,7 @@
 import gpuhunt
 from pydantic import parse_obj_as
 
-import dstack._internal.core.models.resources as resources
-from dstack._internal.cli.services.args import cpu_spec, disk_spec, gpu_spec, port_mapping
+from dstack._internal.cli.services.args import port_mapping
 from dstack._internal.cli.services.configurators.base import (
     ApplyEnvVarsConfiguratorMixin,
     BaseApplyConfigurator,
@@ -26,6 +25,7 @@
     is_git_repo_url,
     register_init_repo_args,
 )
+from dstack._internal.cli.services.resources import apply_resources_args, register_resources_args
 from dstack._internal.cli.utils.common import confirm_ask, console
 from dstack._internal.cli.utils.rich import MultiItemStatus
 from dstack._internal.cli.utils.run import get_runs_table, print_run_plan
@@ -309,29 +309,7 @@ def register_args(cls, parser: argparse.ArgumentParser):
             default=3,
         )
         cls.register_env_args(configuration_group)
-        configuration_group.add_argument(
-            "--cpu",
-            type=cpu_spec,
-            help="Request CPU for the run. "
-            "The format is [code]ARCH[/]:[code]COUNT[/] (all parts are optional)",
-            dest="cpu_spec",
-            metavar="SPEC",
-        )
-        configuration_group.add_argument(
-            "--gpu",
-            type=gpu_spec,
-            help="Request GPU for the run. "
-            "The format is [code]NAME[/]:[code]COUNT[/]:[code]MEMORY[/] (all parts are optional)",
-            dest="gpu_spec",
-            metavar="SPEC",
-        )
-        configuration_group.add_argument(
-            "--disk",
-            type=disk_spec,
-            help="Request the size range of disk for the run. Example [code]--disk 100GB..[/].",
-            metavar="RANGE",
-            dest="disk_spec",
-        )
+        register_resources_args(configuration_group)
         register_profile_args(parser)
         repo_group = parser.add_argument_group("Repo Options")
         repo_group.add_argument(
@@ -359,16 +337,10 @@ def register_args(cls, parser: argparse.ArgumentParser):
         register_init_repo_args(repo_group)
 
     def apply_args(self, conf: RunConfigurationT, args: argparse.Namespace):
+        apply_resources_args(args, conf)
         apply_profile_args(args, conf)
         if args.run_name:
             conf.name = args.run_name
-        if args.cpu_spec:
-            conf.resources.cpu = resources.CPUSpec.parse_obj(args.cpu_spec)
-        if args.gpu_spec:
-            conf.resources.gpu = resources.GPUSpec.parse_obj(args.gpu_spec)
-        if args.disk_spec:
-            conf.resources.disk = args.disk_spec
-
         self.apply_env_vars(conf.env, args)
         self.interpolate_env(conf)
 

src/dstack/_internal/cli/services/resources.py

@@ -0,0 +1,54 @@
+import argparse
+
+from dstack._internal.cli.services.args import cpu_spec, disk_spec, gpu_spec, memory_spec
+from dstack._internal.cli.services.configurators.base import ArgsParser
+from dstack._internal.core.models import resources
+from dstack._internal.core.models.configurations import AnyRunConfiguration
+
+
+def register_resources_args(parser: ArgsParser) -> None:
+    parser.add_argument(
+        "--cpu",
+        type=cpu_spec,
+        help=(
+            "Request CPU for the run."
+            " The format is [code]ARCH[/]:[code]COUNT[/] (all parts are optional)"
+        ),
+        dest="cpu_spec",
+        metavar="SPEC",
+    )
+    parser.add_argument(
+        "--gpu",
+        type=gpu_spec,
+        help=(
+            "Request GPU for the run."
+            " The format is [code]NAME[/]:[code]COUNT[/]:[code]MEMORY[/] (all parts are optional)"
+        ),
+        dest="gpu_spec",
+        metavar="SPEC",
+    )
+    parser.add_argument(
+        "--memory",
+        type=memory_spec,
+        help="Request the size range of RAM for the run. Example [code]--memory 128GB..256GB[/]",
+        dest="memory_spec",
+        metavar="RANGE",
+    )
+    parser.add_argument(
+        "--disk",
+        type=disk_spec,
+        help="Request the size range of disk for the run. Example [code]--disk 100GB..[/]",
+        dest="disk_spec",
+        metavar="RANGE",
+    )
+
+
+def apply_resources_args(args: argparse.Namespace, conf: AnyRunConfiguration) -> None:
+    if args.cpu_spec:
+        conf.resources.cpu = resources.CPUSpec.parse_obj(args.cpu_spec)
+    if args.gpu_spec:
+        conf.resources.gpu = resources.GPUSpec.parse_obj(args.gpu_spec)
+    if args.memory_spec:
+        conf.resources.memory = args.memory_spec
+    if args.disk_spec:
+        conf.resources.disk = args.disk_spec