diff --git a/src/dstack/_internal/cli/commands/offer.py b/src/dstack/_internal/cli/commands/offer.py index bc6bb0a5d..0e4be1d5c 100644 --- a/src/dstack/_internal/cli/commands/offer.py +++ b/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) diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index 3d126dd34..fc76fe43e 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/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) diff --git a/src/dstack/_internal/cli/services/resources.py b/src/dstack/_internal/cli/services/resources.py new file mode 100644 index 000000000..e81b6078d --- /dev/null +++ b/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