Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 2 additions & 24 deletions src/dstack/_internal/cli/commands/offer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)


Expand Down
36 changes: 4 additions & 32 deletions src/dstack/_internal/cli/services/configurators/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)

Expand Down
54 changes: 54 additions & 0 deletions src/dstack/_internal/cli/services/resources.py
Original file line number Diff line number Diff line change
@@ -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