dstack 0.19.27__py3-none-any.whl → 0.19.28__py3-none-any.whl

dstack/_internal/cli/commands/__init__.py

@@ -1,20 +1,22 @@
 import argparse
 import os
+import shlex
 from abc import ABC, abstractmethod
-from typing import List, Optional
+from typing import ClassVar, Optional
 
 from rich_argparse import RichHelpFormatter
 
 from dstack._internal.cli.services.completion import ProjectNameCompleter
-from dstack._internal.cli.utils.common import configure_logging
+from dstack._internal.core.errors import CLIError
 from dstack.api import Client
 
 
 class BaseCommand(ABC):
-    NAME: str = "name the command"
-    DESCRIPTION: str = "describe the command"
-    DEFAULT_HELP: bool = True
-    ALIASES: Optional[List[str]] = None
+    NAME: ClassVar[str] = "name the command"
+    DESCRIPTION: ClassVar[str] = "describe the command"
+    DEFAULT_HELP: ClassVar[bool] = True
+    ALIASES: ClassVar[Optional[list[str]]] = None
+    ACCEPT_EXTRA_ARGS: ClassVar[bool] = False
 
     def __init__(self, parser: argparse.ArgumentParser):
         self._parser = parser
@@ -50,7 +52,8 @@ class BaseCommand(ABC):
 
     @abstractmethod
     def _command(self, args: argparse.Namespace):
-        pass
+        if not self.ACCEPT_EXTRA_ARGS and args.extra_args:
+            raise CLIError(f"Unrecognized arguments: {shlex.join(args.extra_args)}")
 
 
 class APIBaseCommand(BaseCommand):
@@ -65,5 +68,5 @@ class APIBaseCommand(BaseCommand):
         ).completer = ProjectNameCompleter()  # type: ignore[attr-defined]
 
     def _command(self, args: argparse.Namespace):
-        configure_logging()
+        super()._command(args)
         self.api = Client.from_config(project_name=args.project)

dstack/_internal/cli/commands/apply.py

@@ -1,4 +1,5 @@
 import argparse
+import shlex
 
 from argcomplete import FilesCompleter  # type: ignore[attr-defined]
 
@@ -19,6 +20,7 @@ class ApplyCommand(APIBaseCommand):
     NAME = "apply"
     DESCRIPTION = "Apply a configuration"
     DEFAULT_HELP = False
+    ACCEPT_EXTRA_ARGS = True
 
     def _register(self):
         super()._register()
@@ -84,13 +86,14 @@ class ApplyCommand(APIBaseCommand):
             configurator_class = get_apply_configurator_class(configuration.type)
             configurator = configurator_class(api_client=self.api)
             configurator_parser = configurator.get_parser()
-            known, unknown = configurator_parser.parse_known_args(args.unknown)
+            configurator_args, unknown_args = configurator_parser.parse_known_args(args.extra_args)
+            if unknown_args:
+                raise CLIError(f"Unrecognized arguments: {shlex.join(unknown_args)}")
             configurator.apply_configuration(
                 conf=configuration,
                 configuration_path=configuration_path,
                 command_args=args,
-                configurator_args=known,
-                unknown_args=unknown,
+                configurator_args=configurator_args,
             )
         except KeyboardInterrupt:
             console.print("\nOperation interrupted by user. Exiting...")

dstack/_internal/cli/commands/completion.py

@@ -1,3 +1,5 @@
+import argparse
+
 import argcomplete
 
 from dstack._internal.cli.commands import BaseCommand
@@ -15,6 +17,6 @@ class CompletionCommand(BaseCommand):
             choices=["bash", "zsh"],
         )
 
-    def _command(self, args):
+    def _command(self, args: argparse.Namespace):
         super()._command(args)
         print(argcomplete.shellcode(["dstack"], shell=args.shell))  # type: ignore[attr-defined]

dstack/_internal/cli/commands/config.py

@@ -40,6 +40,7 @@ class ConfigCommand(BaseCommand):
         )
 
     def _command(self, args: argparse.Namespace):
+        super()._command(args)
         config_manager = ConfigManager()
         if args.remove:
             config_manager.delete_project(args.project)

dstack/_internal/cli/commands/init.py

@@ -9,7 +9,7 @@ from dstack._internal.cli.services.repos import (
     is_git_repo_url,
     register_init_repo_args,
 )
-from dstack._internal.cli.utils.common import configure_logging, confirm_ask, console, warn
+from dstack._internal.cli.utils.common import confirm_ask, console, warn
 from dstack._internal.core.errors import ConfigurationError
 from dstack._internal.core.models.repos.remote import RemoteRepo
 from dstack._internal.core.services.configs import ConfigManager
@@ -52,7 +52,7 @@ class InitCommand(BaseCommand):
         )
 
     def _command(self, args: argparse.Namespace):
-        configure_logging()
+        super()._command(args)
 
         repo_path: Optional[Path] = None
         repo_url: Optional[str] = None

dstack/_internal/cli/commands/offer.py

@@ -99,7 +99,7 @@ class OfferCommand(APIBaseCommand):
         conf = TaskConfiguration(commands=[":"])
 
         configurator = OfferConfigurator(api_client=self.api)
-        configurator.apply_args(conf, args, [])
+        configurator.apply_args(conf, args)
         profile = load_profile(Path.cwd(), profile_name=args.profile)
 
         run_spec = RunSpec(

dstack/_internal/cli/commands/project.py

@@ -67,6 +67,7 @@ class ProjectCommand(BaseCommand):
         set_default_parser.set_defaults(subfunc=self._set_default)
 
     def _command(self, args: argparse.Namespace):
+        super()._command(args)
         if not hasattr(args, "subfunc"):
             args.subfunc = self._list
         args.subfunc(args)

dstack/_internal/cli/commands/server.py

@@ -1,5 +1,5 @@
+import argparse
 import os
-from argparse import Namespace
 
 from dstack._internal import settings
 from dstack._internal.cli.commands import BaseCommand
@@ -53,7 +53,7 @@ class ServerCommand(BaseCommand):
         )
         self._parser.add_argument("--token", type=str, help="The admin user token")
 
-    def _command(self, args: Namespace):
+    def _command(self, args: argparse.Namespace):
         super()._command(args)
 
         if not UVICORN_INSTALLED:

dstack/_internal/cli/main.py

@@ -83,7 +83,7 @@ def main():
     argcomplete.autocomplete(parser, always_complete_options=False)
 
     args, unknown_args = parser.parse_known_args()
-    args.unknown = unknown_args
+    args.extra_args = unknown_args
 
     try:
         check_for_updates()

dstack/_internal/cli/services/configurators/base.py

@@ -1,7 +1,7 @@
 import argparse
 import os
 from abc import ABC, abstractmethod
-from typing import Generic, List, TypeVar, Union, cast
+from typing import ClassVar, Generic, List, TypeVar, Union, cast
 
 from dstack._internal.cli.services.args import env_var
 from dstack._internal.core.errors import ConfigurationError
@@ -18,7 +18,7 @@ ApplyConfigurationT = TypeVar("ApplyConfigurationT", bound=AnyApplyConfiguration
 
 
 class BaseApplyConfigurator(ABC, Generic[ApplyConfigurationT]):
-    TYPE: ApplyConfigurationType
+    TYPE: ClassVar[ApplyConfigurationType]
 
     def __init__(self, api_client: Client):
         self.api = api_client
@@ -30,7 +30,6 @@ class BaseApplyConfigurator(ABC, Generic[ApplyConfigurationT]):
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
         """
         Implements `dstack apply` for a given configuration type.
@@ -40,7 +39,6 @@ class BaseApplyConfigurator(ABC, Generic[ApplyConfigurationT]):
             configuration_path: The path to the configuration file.
             command_args: The args parsed by `dstack apply`.
             configurator_args: The known args parsed by `cls.get_parser()`.
-            unknown_args: The unknown args after parsing by `cls.get_parser()`.
         """
         pass
 

dstack/_internal/cli/services/configurators/fleet.py

@@ -1,7 +1,7 @@
 import argparse
 import time
 from pathlib import Path
-from typing import List, Optional
+from typing import Optional
 
 from rich.table import Table
 
@@ -46,7 +46,7 @@ logger = get_logger(__name__)
 
 
 class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator[FleetConfiguration]):
-    TYPE: ApplyConfigurationType = ApplyConfigurationType.FLEET
+    TYPE = ApplyConfigurationType.FLEET
 
     def apply_configuration(
         self,
@@ -54,9 +54,8 @@ class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator[Fle
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
-        self.apply_args(conf, configurator_args, unknown_args)
+        self.apply_args(conf, configurator_args)
         profile = load_profile(Path.cwd(), None)
         spec = FleetSpec(
             configuration=conf,
@@ -309,7 +308,7 @@ class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator[Fle
         )
         cls.register_env_args(configuration_group)
 
-    def apply_args(self, conf: FleetConfiguration, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: FleetConfiguration, args: argparse.Namespace):
         if args.name:
             conf.name = args.name
         self.apply_env_vars(conf.env, args)

dstack/_internal/cli/services/configurators/gateway.py

@@ -1,6 +1,5 @@
 import argparse
 import time
-from typing import List
 
 from rich.table import Table
 
@@ -27,7 +26,7 @@ from dstack.api._public import Client
 
 
 class GatewayConfigurator(BaseApplyConfigurator[GatewayConfiguration]):
-    TYPE: ApplyConfigurationType = ApplyConfigurationType.GATEWAY
+    TYPE = ApplyConfigurationType.GATEWAY
 
     def apply_configuration(
         self,
@@ -35,9 +34,8 @@ class GatewayConfigurator(BaseApplyConfigurator[GatewayConfiguration]):
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
-        self.apply_args(conf, configurator_args, unknown_args)
+        self.apply_args(conf, configurator_args)
         spec = GatewaySpec(
             configuration=conf,
             configuration_path=configuration_path,
@@ -179,7 +177,7 @@ class GatewayConfigurator(BaseApplyConfigurator[GatewayConfiguration]):
             help="The gateway name",
         )
 
-    def apply_args(self, conf: GatewayConfiguration, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: GatewayConfiguration, args: argparse.Namespace):
         if args.name:
             conf.name = args.name
 

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

@@ -1,4 +1,5 @@
 import argparse
+import shlex
 import subprocess
 import sys
 import time
@@ -35,6 +36,7 @@ from dstack._internal.core.models.configurations import (
     LEGACY_REPO_DIR,
     AnyRunConfiguration,
     ApplyConfigurationType,
+    ConfigurationWithCommandsParams,
     ConfigurationWithPortsParams,
     DevEnvironmentConfiguration,
     PortMapping,
@@ -80,20 +82,17 @@ class BaseRunConfigurator(
     ApplyEnvVarsConfiguratorMixin,
     BaseApplyConfigurator[RunConfigurationT],
 ):
-    TYPE: ApplyConfigurationType
-
     def apply_configuration(
         self,
         conf: RunConfigurationT,
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
         if configurator_args.repo and configurator_args.no_repo:
             raise CLIError("Either --repo or --no-repo can be specified")
 
-        self.apply_args(conf, configurator_args, unknown_args)
+        self.apply_args(conf, configurator_args)
         self.validate_gpu_vendor_and_image(conf)
         self.validate_cpu_arch_and_image(conf)
 
@@ -395,7 +394,7 @@ class BaseRunConfigurator(
         )
         register_init_repo_args(repo_group)
 
-    def apply_args(self, conf: RunConfigurationT, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: RunConfigurationT, args: argparse.Namespace):
         apply_profile_args(args, conf)
         if args.run_name:
             conf.name = args.run_name
@@ -408,16 +407,6 @@ class BaseRunConfigurator(
 
         self.apply_env_vars(conf.env, args)
         self.interpolate_env(conf)
-        self.interpolate_run_args(conf.setup, unknown)
-
-    def interpolate_run_args(self, value: List[str], unknown):
-        run_args = " ".join(unknown)
-        interpolator = VariablesInterpolator({"run": {"args": run_args}}, skip=["secrets"])
-        try:
-            for i in range(len(value)):
-                value[i] = interpolator.interpolate_or_error(value[i])
-        except InterpolatorError as e:
-            raise ConfigurationError(e.args[0])
 
     def interpolate_env(self, conf: RunConfigurationT):
         env_dict = conf.env.as_dict()
@@ -701,18 +690,50 @@ class RunWithPortsConfiguratorMixin:
             conf.ports = list(_merge_ports(conf.ports, args.ports).values())
 
 
-class TaskConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigurator):
+class RunWithCommandsConfiguratorMixin:
+    @classmethod
+    def register_commands_args(cls, parser: argparse.ArgumentParser):
+        parser.add_argument(
+            "run_args",
+            help=(
+                "Run arguments. Available in the configuration [code]commands[/code] as"
+                " [code]${{ run.args }}[/code]."
+                " Use [code]--[/code] to separate run options from [code]dstack[/code] options"
+            ),
+            nargs="*",
+            metavar="RUN_ARGS",
+        )
+
+    def apply_commands_args(
+        self,
+        conf: ConfigurationWithCommandsParams,
+        args: argparse.Namespace,
+    ):
+        commands = conf.commands
+        run_args = shlex.join(args.run_args)
+        interpolator = VariablesInterpolator({"run": {"args": run_args}}, skip=["secrets"])
+        try:
+            for i, command in enumerate(commands):
+                commands[i] = interpolator.interpolate_or_error(command)
+        except InterpolatorError as e:
+            raise ConfigurationError(e.args[0])
+
+
+class TaskConfigurator(
+    RunWithPortsConfiguratorMixin, RunWithCommandsConfiguratorMixin, BaseRunConfigurator
+):
     TYPE = ApplyConfigurationType.TASK
 
     @classmethod
     def register_args(cls, parser: argparse.ArgumentParser):
         super().register_args(parser)
         cls.register_ports_args(parser)
+        cls.register_commands_args(parser)
 
-    def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace, unknown: List[str]):
-        super().apply_args(conf, args, unknown)
+    def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace):
+        super().apply_args(conf, args)
         self.apply_ports_args(conf, args)
-        self.interpolate_run_args(conf.commands, unknown)
+        self.apply_commands_args(conf, args)
 
 
 class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigurator):
@@ -723,10 +744,8 @@ class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigura
         super().register_args(parser)
         cls.register_ports_args(parser)
 
-    def apply_args(
-        self, conf: DevEnvironmentConfiguration, args: argparse.Namespace, unknown: List[str]
-    ):
-        super().apply_args(conf, args, unknown)
+    def apply_args(self, conf: DevEnvironmentConfiguration, args: argparse.Namespace):
+        super().apply_args(conf, args)
         self.apply_ports_args(conf, args)
         if conf.ide == "vscode" and conf.version is None:
             conf.version = _detect_vscode_version()
@@ -746,12 +765,17 @@ class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigura
                 )
 
 
-class ServiceConfigurator(BaseRunConfigurator):
+class ServiceConfigurator(RunWithCommandsConfiguratorMixin, BaseRunConfigurator):
     TYPE = ApplyConfigurationType.SERVICE
 
-    def apply_args(self, conf: ServiceConfiguration, args: argparse.Namespace, unknown: List[str]):
-        super().apply_args(conf, args, unknown)
-        self.interpolate_run_args(conf.commands, unknown)
+    @classmethod
+    def register_args(cls, parser: argparse.ArgumentParser):
+        super().register_args(parser)
+        cls.register_commands_args(parser)
+
+    def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace):
+        super().apply_args(conf, args)
+        self.apply_commands_args(conf, args)
 
 
 def _merge_ports(conf: List[PortMapping], args: List[PortMapping]) -> Dict[int, PortMapping]:

dstack/_internal/cli/services/configurators/volume.py

@@ -1,6 +1,5 @@
 import argparse
 import time
-from typing import List
 
 from rich.table import Table
 
@@ -26,7 +25,7 @@ from dstack.api._public import Client
 
 
 class VolumeConfigurator(BaseApplyConfigurator[VolumeConfiguration]):
-    TYPE: ApplyConfigurationType = ApplyConfigurationType.VOLUME
+    TYPE = ApplyConfigurationType.VOLUME
 
     def apply_configuration(
         self,
@@ -34,9 +33,8 @@ class VolumeConfigurator(BaseApplyConfigurator[VolumeConfiguration]):
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
-        self.apply_args(conf, configurator_args, unknown_args)
+        self.apply_args(conf, configurator_args)
         spec = VolumeSpec(
             configuration=conf,
             configuration_path=configuration_path,
@@ -167,7 +165,7 @@ class VolumeConfigurator(BaseApplyConfigurator[VolumeConfiguration]):
             help="The volume name",
         )
 
-    def apply_args(self, conf: VolumeConfiguration, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: VolumeConfiguration, args: argparse.Namespace):
         if args.name:
             conf.name = args.name
 

dstack/_internal/core/compatibility/runs.py

@@ -31,6 +31,8 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeD
         current_resource_excludes["status_message"] = True
         if current_resource.deployment_num == 0:
             current_resource_excludes["deployment_num"] = True
+        if current_resource.fleet is None:
+            current_resource_excludes["fleet"] = True
         apply_plan_excludes["current_resource"] = current_resource_excludes
         current_resource_excludes["run_spec"] = get_run_spec_excludes(current_resource.run_spec)
         job_submissions_excludes: IncludeExcludeDictType = {}

dstack/_internal/core/models/common.py

@@ -1,10 +1,10 @@
 import re
 from enum import Enum
-from typing import Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Mapping, Optional, Union
 
 import orjson
 from pydantic import Field
-from pydantic_duality import DualBaseModel
+from pydantic_duality import generate_dual_base_model
 from typing_extensions import Annotated
 
 from dstack._internal.utils.json_utils import pydantic_orjson_dumps
@@ -17,46 +17,73 @@ IncludeExcludeDictType = dict[
 IncludeExcludeType = Union[IncludeExcludeSetType, IncludeExcludeDictType]
 
 
+class CoreConfig:
+    json_loads = orjson.loads
+    json_dumps = pydantic_orjson_dumps
+
+
+# All dstack models inherit from pydantic-duality's DualBaseModel.
 # DualBaseModel creates two classes for the model:
 # one with extra = "forbid" (CoreModel/CoreModel.__request__),
 # and another with extra = "ignore" (CoreModel.__response__).
-# This allows to use the same model both for a strict parsing of the user input and
-# for a permissive parsing of the server responses.
-class CoreModel(DualBaseModel):
-    class Config:
-        json_loads = orjson.loads
-        json_dumps = pydantic_orjson_dumps
-
-    def json(
-        self,
-        *,
-        include: Optional[IncludeExcludeType] = None,
-        exclude: Optional[IncludeExcludeType] = None,
-        by_alias: bool = False,
-        skip_defaults: Optional[bool] = None,  # ignore as it's deprecated
-        exclude_unset: bool = False,
-        exclude_defaults: bool = False,
-        exclude_none: bool = False,
-        encoder: Optional[Callable[[Any], Any]] = None,
-        models_as_dict: bool = True,  # does not seems to be needed by dstack or dependencies
-        **dumps_kwargs: Any,
-    ) -> str:
-        """
-        Override `json()` method so that it calls `dict()`.
-        Allows changing how models are serialized by overriding `dict()` only.
-        By default, `json()` won't call `dict()`, so changes applied in `dict()` won't take place.
-        """
-        data = self.dict(
-            by_alias=by_alias,
-            include=include,
-            exclude=exclude,
-            exclude_unset=exclude_unset,
-            exclude_defaults=exclude_defaults,
-            exclude_none=exclude_none,
-        )
-        if self.__custom_root_type__:
-            data = data["__root__"]
-        return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)
+# This allows to use the same model both for strict parsing of the user input and
+# for permissive parsing of the server responses.
+#
+# We define a func to generate CoreModel dynamically that can be used
+# to define custom Config for both __request__ and __response__ models.
+# Note: Defining config in the model class directly overrides
+# pydantic-duality's base config, breaking __response__.
+def generate_dual_core_model(
+    custom_config: Union[type, Mapping],
+) -> "type[CoreModel]":
+    class CoreModel(generate_dual_base_model(custom_config)):
+        def json(
+            self,
+            *,
+            include: Optional[IncludeExcludeType] = None,
+            exclude: Optional[IncludeExcludeType] = None,
+            by_alias: bool = False,
+            skip_defaults: Optional[bool] = None,  # ignore as it's deprecated
+            exclude_unset: bool = False,
+            exclude_defaults: bool = False,
+            exclude_none: bool = False,
+            encoder: Optional[Callable[[Any], Any]] = None,
+            models_as_dict: bool = True,  # does not seems to be needed by dstack or dependencies
+            **dumps_kwargs: Any,
+        ) -> str:
+            """
+            Override `json()` method so that it calls `dict()`.
+            Allows changing how models are serialized by overriding `dict()` only.
+            By default, `json()` won't call `dict()`, so changes applied in `dict()` won't take place.
+            """
+            data = self.dict(
+                by_alias=by_alias,
+                include=include,
+                exclude=exclude,
+                exclude_unset=exclude_unset,
+                exclude_defaults=exclude_defaults,
+                exclude_none=exclude_none,
+            )
+            if self.__custom_root_type__:
+                data = data["__root__"]
+            return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)
+
+    return CoreModel
+
+
+if TYPE_CHECKING:
+
+    class CoreModel(generate_dual_base_model(CoreConfig)):
+        pass
+else:
+    CoreModel = generate_dual_core_model(CoreConfig)
+
+
+class FrozenConfig(CoreConfig):
+    frozen = True
+
+
+FrozenCoreModel = generate_dual_core_model(FrozenConfig)
 
 
 class Duration(int):
@@ -93,7 +120,7 @@ class Duration(int):
         raise ValueError(f"Cannot parse the duration {v}")
 
 
-class RegistryAuth(CoreModel):
+class RegistryAuth(FrozenCoreModel):
     """
     Credentials for pulling a private Docker image.
 
@@ -105,9 +132,6 @@ class RegistryAuth(CoreModel):
     username: Annotated[str, Field(description="The username")]
     password: Annotated[str, Field(description="The password or access token")]
 
-    class Config(CoreModel.Config):
-        frozen = True
-
 
 class ApplyAction(str, Enum):
     CREATE = "create"  # resource is to be created or overridden