Fix typing issues and add pyright to CI (#3011)

* Fix possible unbound vars in background tasks

* Move class Config after fields

* Fix unbound vars in app.py

* Assert provisioning gateways have compute

* Assert region set for ssh instances

* Assert region service configurations

* Fix logging type annotations

* Fix redundant add_project_members line

* Fix Lockset annotations

* Assert configuration types in job configurators

* Fix unbound vars in process_running_jobs

* Check backend is not None

* Assert service and jpd in register_service

* Assert jpd in container_ssh_tunnel

* Assert jpd in ServerProxyRepo

* Fix configuration.model type annotation

* Assert conf.replicas

* Fix abstract AsyncGenerator def

* Fix ProbeConfig type annotations

* Fix max_duration and stop_duration type annotations

* Fix idle_duration type annotations

* Fix volumes and files type annotations

* Fix retry.duration type annotation

* Do not define Storage implementations when deps missing

* Fix gateway domain None

* Overload get_backend_config

* Do not define LogStorage implementations when deps missing

* Fix filelog typing

* Use async_sessionmaker

* Assert proxy_jump.ssh_key

* Ignore type errors from deps

* Forbid entrypoint for dev-environment

Fixes #3002

* Fix vscode and cursor __init__ annotations

* Pass probe_spec.body as content

* Assert gateway configuration.name

* Fix unbound next_token

* Cast path to str

* Fix unbound success var

* Fix typing

* Add pyright config

* Run pyright in CI

* Run pyright as part of tests

* Ignore type for entry_points

* Fix _detect_vscode_version

* Remove run_name assert

* Fix add_extra_schema_types for one $ref

* Replace ConfigurationWith extensions with mixins

* Fix unbound spec_json

* Remove huggingface api

* Type check plugins

* Fix BaseApplyConfigurator generics

* Type check core/services

* Fix services.gpus typing

* Document pyright in Contributing

Author: r4victor

.github/workflows/build-artifacts.yml

@@ -73,6 +73,10 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: uv sync --all-extras
+      - name: Run pyright
+        uses: jakebailey/pyright-action@v2
+        with:
+          pylance-version: latest-release
       - name: Download frontend build
         uses: actions/download-artifact@v4
         with:

contributing/DEVELOPMENT.md

@@ -25,12 +25,26 @@ uv sync --all-extras
 
 Alternatively, if you want to manage virtual environments by yourself, you can install `dstack` into the activated virtual environment with `uv sync --all-extras --active`.
 
-## 4. (Recommended) Install pre-commits:
+## 4. (Recommended) Install pre-commit hooks:
+
+Code formatting and linting can be done automatically on each commit with `pre-commit` hooks:
 
 ```shell
 uv run pre-commit install
 ```
 
-## 5. Frontend
+## 5. (Recommended) Use pyright:
+
+The CI runs `pyright` for type checking `dstack` Python code.
+So we recommend you configure your IDE to use `pyright`/`pylance` with `standard` type checking mode.
+
+You can also install `pyright` and run it from the CLI:
+
+```shell
+uv tool install pyright
+pyright -p .
+```
+
+## 6. Frontend
 
 See [FRONTEND.md](FRONTEND.md) for the details on how to build and develop the frontend.

pyproject.toml

@@ -78,6 +78,17 @@ pattern = '<picture>\s*|<source[^>]*>\s*|\s*</picture>|<video[^>]*>\s*|</video>\
 replacement = ''
 ignore-case = true
 
+[tool.pyright]
+include = [
+    "src/dstack/plugins",
+    "src/dstack/_internal/server",
+    "src/dstack/_internal/core/services",
+    "src/dstack/_internal/cli/services/configurators",
+]
+ignore = [
+    "src/dstack/_internal/server/migrations/versions",
+]
+
 [dependency-groups]
 dev = [
     "httpx>=0.28.1",

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

@@ -3,7 +3,11 @@
 from typing import List
 
 from dstack._internal.cli.commands import APIBaseCommand
-from dstack._internal.cli.services.configurators.run import BaseRunConfigurator
+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.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
@@ -18,11 +22,8 @@ class OfferConfigurator(BaseRunConfigurator):
     TYPE = ApplyConfigurationType.TASK
 
     @classmethod
-    def register_args(
-        cls,
-        parser: argparse.ArgumentParser,
-    ):
-        super().register_args(parser, default_max_offers=50)
+    def register_args(cls, parser: argparse.ArgumentParser):
+        configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options")
         parser.add_argument(
             "--group-by",
             action="append",
@@ -33,6 +34,43 @@ def register_args(
                 "Can be repeated or comma-separated (e.g. [code]--group-by gpu,backend[/code])."
             ),
         )
+        configuration_group.add_argument(
+            "-n",
+            "--name",
+            dest="run_name",
+            help="The name of the run. If not specified, a random name is assigned",
+        )
+        configuration_group.add_argument(
+            "--max-offers",
+            help="Number of offers to show in the run plan",
+            type=int,
+            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_profile_args(parser)
 
 
 class OfferCommand(APIBaseCommand):
@@ -117,7 +155,7 @@ def _process_group_by_args(self, group_by_args: List[str]) -> List[str]:
 
         return processed
 
-    def _list_gpus(self, args: List[str], run_spec: RunSpec) -> List[GpuGroup]:
+    def _list_gpus(self, args: argparse.Namespace, run_spec: RunSpec) -> List[GpuGroup]:
         group_by = [g for g in args.group_by if g != "gpu"] or None
         return self.api.client.gpus.list_gpus(
             self.api.project,

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

@@ -24,7 +24,9 @@
 APPLY_STDIN_NAME = "-"
 
 
-apply_configurators_mapping: Dict[ApplyConfigurationType, Type[BaseApplyConfigurator]] = {
+apply_configurators_mapping: Dict[
+    ApplyConfigurationType, Type[BaseApplyConfigurator[AnyApplyConfiguration]]
+] = {
     cls.TYPE: cls
     for cls in [
         DevEnvironmentConfigurator,
@@ -47,7 +49,9 @@
 }
 
 
-def get_apply_configurator_class(configurator_type: str) -> Type[BaseApplyConfigurator]:
+def get_apply_configurator_class(
+    configurator_type: str,
+) -> Type[BaseApplyConfigurator[AnyApplyConfiguration]]:
     return apply_configurators_mapping[ApplyConfigurationType(configurator_type)]
 
 

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

@@ -1,7 +1,7 @@
 import argparse
 import os
 from abc import ABC, abstractmethod
-from typing import List, Optional, Union, cast
+from typing import Generic, List, Optional, TypeVar, Union, cast
 
 from dstack._internal.cli.services.args import env_var
 from dstack._internal.core.errors import ConfigurationError
@@ -15,8 +15,10 @@
 
 ArgsParser = Union[argparse._ArgumentGroup, argparse.ArgumentParser]
 
+ApplyConfigurationT = TypeVar("ApplyConfigurationT", bound=AnyApplyConfiguration)
 
-class BaseApplyConfigurator(ABC):
+
+class BaseApplyConfigurator(ABC, Generic[ApplyConfigurationT]):
     TYPE: ApplyConfigurationType
 
     def __init__(self, api_client: Client):
@@ -25,7 +27,7 @@ def __init__(self, api_client: Client):
     @abstractmethod
     def apply_configuration(
         self,
-        conf: AnyApplyConfiguration,
+        conf: ApplyConfigurationT,
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
@@ -48,7 +50,7 @@ def apply_configuration(
     @abstractmethod
     def delete_configuration(
         self,
-        conf: AnyApplyConfiguration,
+        conf: ApplyConfigurationT,
         configuration_path: str,
         command_args: argparse.Namespace,
     ):

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

@@ -46,7 +46,7 @@
 logger = get_logger(__name__)
 
 
-class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
+class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator[FleetConfiguration]):
     TYPE: ApplyConfigurationType = ApplyConfigurationType.FLEET
 
     def apply_configuration(

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

@@ -27,7 +27,7 @@
 from dstack.api._public import Client
 
 
-class GatewayConfigurator(BaseApplyConfigurator):
+class GatewayConfigurator(BaseApplyConfigurator[GatewayConfiguration]):
     TYPE: ApplyConfigurationType = ApplyConfigurationType.GATEWAY
 
     def apply_configuration(

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

@@ -3,7 +3,7 @@
 import sys
 import time
 from pathlib import Path
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, TypeVar
 
 import gpuhunt
 from pydantic import parse_obj_as
@@ -33,8 +33,7 @@
 from dstack._internal.core.models.configurations import (
     AnyRunConfiguration,
     ApplyConfigurationType,
-    BaseRunConfiguration,
-    BaseRunConfigurationWithPorts,
+    ConfigurationWithPortsParams,
     DevEnvironmentConfiguration,
     PortMapping,
     RunConfigurationType,
@@ -63,13 +62,18 @@
 
 logger = get_logger(__name__)
 
+RunConfigurationT = TypeVar("RunConfigurationT", bound=AnyRunConfiguration)
 
-class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
+
+class BaseRunConfigurator(
+    ApplyEnvVarsConfiguratorMixin,
+    BaseApplyConfigurator[RunConfigurationT],
+):
     TYPE: ApplyConfigurationType
 
     def apply_configuration(
         self,
-        conf: BaseRunConfiguration,
+        conf: RunConfigurationT,
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
@@ -267,7 +271,7 @@ def apply_configuration(
 
     def delete_configuration(
         self,
-        conf: AnyRunConfiguration,
+        conf: RunConfigurationT,
         configuration_path: str,
         command_args: argparse.Namespace,
     ):
@@ -293,7 +297,7 @@ def delete_configuration(
         console.print(f"Run [code]{conf.name}[/] deleted")
 
     @classmethod
-    def register_args(cls, parser: argparse.ArgumentParser, default_max_offers: int = 3):
+    def register_args(cls, parser: argparse.ArgumentParser):
         configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options")
         configuration_group.add_argument(
             "-n",
@@ -305,7 +309,7 @@ def register_args(cls, parser: argparse.ArgumentParser, default_max_offers: int
             "--max-offers",
             help="Number of offers to show in the run plan",
             type=int,
-            default=default_max_offers,
+            default=3,
         )
         cls.register_env_args(configuration_group)
         configuration_group.add_argument(
@@ -333,7 +337,7 @@ def register_args(cls, parser: argparse.ArgumentParser, default_max_offers: int
         )
         register_profile_args(parser)
 
-    def apply_args(self, conf: BaseRunConfiguration, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: RunConfigurationT, args: argparse.Namespace, unknown: List[str]):
         apply_profile_args(args, conf)
         if args.run_name:
             conf.name = args.run_name
@@ -357,7 +361,7 @@ def interpolate_run_args(self, value: List[str], unknown):
         except InterpolatorError as e:
             raise ConfigurationError(e.args[0])
 
-    def interpolate_env(self, conf: BaseRunConfiguration):
+    def interpolate_env(self, conf: RunConfigurationT):
         env_dict = conf.env.as_dict()
         interpolator = VariablesInterpolator({"env": env_dict}, skip=["secrets"])
         try:
@@ -377,7 +381,7 @@ def interpolate_env(self, conf: BaseRunConfiguration):
         except InterpolatorError as e:
             raise ConfigurationError(e.args[0])
 
-    def validate_gpu_vendor_and_image(self, conf: BaseRunConfiguration) -> None:
+    def validate_gpu_vendor_and_image(self, conf: RunConfigurationT) -> None:
         """
         Infers and sets `resources.gpu.vendor` if not set, requires `image` if the vendor is AMD.
         """
@@ -438,7 +442,7 @@ def validate_gpu_vendor_and_image(self, conf: BaseRunConfiguration) -> None:
                 "`image` is required if `resources.gpu.vendor` is `tenstorrent`"
             )
 
-    def validate_cpu_arch_and_image(self, conf: BaseRunConfiguration) -> None:
+    def validate_cpu_arch_and_image(self, conf: RunConfigurationT) -> None:
         """
         Infers `resources.cpu.arch` if not set, requires `image` if the architecture is ARM.
         """
@@ -462,10 +466,9 @@ def validate_cpu_arch_and_image(self, conf: BaseRunConfiguration) -> None:
             raise ConfigurationError("`image` is required if `resources.cpu.arch` is `arm`")
 
 
-class RunWithPortsConfigurator(BaseRunConfigurator):
+class RunWithPortsConfiguratorMixin:
     @classmethod
-    def register_args(cls, parser: argparse.ArgumentParser):
-        super().register_args(parser)
+    def register_ports_args(cls, parser: argparse.ArgumentParser):
         parser.add_argument(
             "-p",
             "--port",
@@ -482,29 +485,42 @@ def register_args(cls, parser: argparse.ArgumentParser):
             metavar="HOST",
         )
 
-    def apply_args(
-        self, conf: BaseRunConfigurationWithPorts, args: argparse.Namespace, unknown: List[str]
+    def apply_ports_args(
+        self,
+        conf: ConfigurationWithPortsParams,
+        args: argparse.Namespace,
     ):
-        super().apply_args(conf, args, unknown)
         if args.ports:
             conf.ports = list(_merge_ports(conf.ports, args.ports).values())
 
 
-class TaskConfigurator(RunWithPortsConfigurator):
+class TaskConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigurator):
     TYPE = ApplyConfigurationType.TASK
 
+    @classmethod
+    def register_args(cls, parser: argparse.ArgumentParser):
+        super().register_args(parser)
+        cls.register_ports_args(parser)
+
     def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace, unknown: List[str]):
         super().apply_args(conf, args, unknown)
+        self.apply_ports_args(conf, args)
         self.interpolate_run_args(conf.commands, unknown)
 
 
-class DevEnvironmentConfigurator(RunWithPortsConfigurator):
+class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigurator):
     TYPE = ApplyConfigurationType.DEV_ENVIRONMENT
 
+    @classmethod
+    def register_args(cls, parser: argparse.ArgumentParser):
+        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)
+        self.apply_ports_args(conf, args)
         if conf.ide == "vscode" and conf.version is None:
             conf.version = _detect_vscode_version()
             if conf.version is None:
@@ -674,6 +690,8 @@ def render_run_spec_diff(old_spec: RunSpec, new_spec: RunSpec) -> Optional[str]:
             if type(old_spec.profile) is not type(new_spec.profile):
                 item = NestedListItem("Profile")
             else:
+                assert old_spec.profile is not None
+                assert new_spec.profile is not None
                 item = NestedListItem(
                     "Profile properties:",
                     children=[

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

@@ -26,7 +26,7 @@
 from dstack.api._public import Client
 
 
-class VolumeConfigurator(BaseApplyConfigurator):
+class VolumeConfigurator(BaseApplyConfigurator[VolumeConfiguration]):
     TYPE: ApplyConfigurationType = ApplyConfigurationType.VOLUME
 
     def apply_configuration(