From e8137f644cbc4864122d778ae07130ee1b000474 Mon Sep 17 00:00:00 2001 From: Egor Sklyarov Date: Tue, 18 Jul 2023 16:16:06 +0400 Subject: [PATCH] Add --ports to `dstack run` (#573) * Add --ports to CLI, align port forwarding syntax with docker * Update --ports in docs --- .../_internal/configurators/__init__.py | 23 +++++--- cli/dstack/_internal/configurators/ports.py | 53 +++++-------------- cli/dstack/_internal/core/configuration.py | 38 ++++++++++++- docs/docs/reference/cli/run.md | 2 +- 4 files changed, 68 insertions(+), 48 deletions(-) diff --git a/cli/dstack/_internal/configurators/__init__.py b/cli/dstack/_internal/configurators/__init__.py index 9448ad033..9ae64488c 100644 --- a/cli/dstack/_internal/configurators/__init__.py +++ b/cli/dstack/_internal/configurators/__init__.py @@ -43,6 +43,10 @@ def get_parser( if parser is None: parser = argparse.ArgumentParser(prog=prog, formatter_class=RichHelpFormatter) + parser.add_argument( + "-p", "--ports", metavar="PORT", type=port_mapping, nargs=argparse.ONE_OR_MORE + ) + spot_group = parser.add_mutually_exclusive_group() spot_group.add_argument( "--spot", action="store_const", dest="spot_policy", const=job.SpotPolicy.SPOT @@ -76,6 +80,9 @@ def get_parser( return parser def apply_args(self, args: argparse.Namespace): + if args.ports is not None: + self.conf.ports = list(ports.merge_ports(self.conf.ports, args.ports).values()) + if args.spot_policy is not None: self.profile.spot_policy = args.spot_policy @@ -212,8 +219,8 @@ def app_specs(self) -> List[job.AppSpec]: for i, pm in enumerate(ports.filter_reserved_ports(self.ports())): specs.append( job.AppSpec( - port=pm.port, - map_to_port=pm.map_to_port, + port=pm.container_port, + map_to_port=pm.local_port, app_name=f"app_{i}", ) ) @@ -226,13 +233,12 @@ def python(self) -> str: return PythonVersion(f"{version_info.major}.{version_info.minor}").value def ports(self) -> Dict[int, ports.PortMapping]: - mapping = [ports.PortMapping(p) for p in self.conf.ports] - ports.unique_ports_constraint([pm.port for pm in mapping]) + ports.unique_ports_constraint([pm.container_port for pm in self.conf.ports]) ports.unique_ports_constraint( - [pm.map_to_port for pm in mapping if pm.map_to_port is not None], + [pm.local_port for pm in self.conf.ports if pm.local_port is not None], error="Mapped port {} is already in use", ) - return {pm.port: pm for pm in mapping} + return {pm.container_port: pm for pm in self.conf.ports} def env(self) -> Dict[str, str]: return self.conf.env @@ -288,3 +294,8 @@ def validate_local_path(path: str, home: Optional[str], working_dir: str) -> str class HomeDirUnsetError(DstackError): pass + + +def port_mapping(v: str) -> ports.PortMapping: + # argparse uses __name__ for handling ValueError + return ports.PortMapping.parse(v) diff --git a/cli/dstack/_internal/configurators/ports.py b/cli/dstack/_internal/configurators/ports.py index c29b4e6d7..467fd1f91 100644 --- a/cli/dstack/_internal/configurators/ports.py +++ b/cli/dstack/_internal/configurators/ports.py @@ -1,7 +1,6 @@ -import argparse -import re -from typing import Dict, Iterator, List, Optional, Union +from typing import Dict, Iterator, List, Optional +from dstack._internal.core.configuration import PortMapping from dstack._internal.core.error import DstackError RESERVED_PORTS_START = 10000 @@ -16,46 +15,20 @@ class PortUsedError(DstackError): pass -class PortMapping: - """ - Valid formats: - - 1234 - - "1234" - - "1234:5678" - """ - - def __init__(self, v: Union[str, int]): - self.port: int - self.map_to_port: Optional[int] = None - - if isinstance(v, int): - self.port = v - return - r = re.search(r"^(\d+)(?::(\d+))?$", v) - if r is None: - raise argparse.ArgumentTypeError(f"{v} is not a valid port or port mapping") - port, map_to_port = r.groups() - self.port = int(port) - if map_to_port is not None: - self.map_to_port = int(map_to_port) - - def __repr__(self): - s = str(self.port) - if self.map_to_port is not None: - s += f":{self.map_to_port}" - return f'PortMapping("{s}")' - - def merge_ports(schema: List[PortMapping], args: List[PortMapping]) -> Dict[int, PortMapping]: - unique_ports_constraint([pm.port for pm in schema], error="Schema port {} is already in use") - unique_ports_constraint([pm.port for pm in args], error="Args port {} is already in use") + unique_ports_constraint( + [pm.container_port for pm in schema], error="Schema port {} is already in use" + ) + unique_ports_constraint( + [pm.container_port for pm in args], error="Args port {} is already in use" + ) - ports = {pm.port: pm for pm in schema} + ports = {pm.container_port: pm for pm in schema} for pm in args: # override schema - ports[pm.port] = pm + ports[pm.container_port] = pm unique_ports_constraint( - [pm.map_to_port for pm in ports.values() if pm.map_to_port is not None], + [pm.local_port for pm in ports.values() if pm.local_port is not None], error="Mapped port {} is already in use", ) return ports @@ -71,12 +44,12 @@ def unique_ports_constraint(ports: List[int], error: str = "Port {} is already i def get_map_to_port(ports: Dict[int, PortMapping], port: int) -> Optional[int]: if port in ports: - return ports[port].map_to_port + return ports[port].local_port return None def filter_reserved_ports(ports: Dict[int, PortMapping]) -> Iterator[PortMapping]: for i in ports.values(): - if RESERVED_PORTS_START <= i.port <= RESERVED_PORTS_END: + if RESERVED_PORTS_START <= i.container_port <= RESERVED_PORTS_END: continue yield i diff --git a/cli/dstack/_internal/core/configuration.py b/cli/dstack/_internal/core/configuration.py index f78a5a7f2..f616dafe4 100644 --- a/cli/dstack/_internal/core/configuration.py +++ b/cli/dstack/_internal/core/configuration.py @@ -1,3 +1,4 @@ +import re from enum import Enum from typing import Dict, List, Optional, Union @@ -5,6 +6,7 @@ from typing_extensions import Annotated, Literal CommandsList = List[str] +ValidPort = conint(gt=0, le=65536) class PythonVersion(str, Enum): @@ -25,6 +27,28 @@ class RegistryAuth(ForbidExtra): password: Annotated[str, Field(description="Password or access token")] +class PortMapping(ForbidExtra): + local_port: Optional[ValidPort] = None + container_port: ValidPort + + @classmethod + def parse(cls, v: str) -> "PortMapping": + """ + Possible values: + - 8080 + - :8080 + - 80:8080 + """ + r = re.search(r"^(?:(\d+)?:)?(\d+)?$", v) + if not r: + raise ValueError(v) + local_port, container_port = r.groups() + return PortMapping( + local_port=None if local_port is None else int(local_port), + container_port=container_port, + ) + + class Artifact(ForbidExtra): path: Annotated[ str, Field(description="The path to the folder that must be stored as an output artifact") @@ -52,7 +76,7 @@ class BaseConfiguration(ForbidExtra): Field(description="The major version of Python\nMutually exclusive with the image"), ] ports: Annotated[ - List[Union[constr(regex=r"^[0-9]+:[0-9]+$"), conint(gt=0, le=65536)]], + List[Union[constr(regex=r"^(?:([0-9]+)?:)?[0-9]+$"), ValidPort, PortMapping]], Field(description="Port numbers/mapping to expose"), ] = [] env: Annotated[ @@ -78,6 +102,18 @@ def convert_python(cls, v, values) -> Optional[PythonVersion]: return PythonVersion(v) return v + @validator("ports") + def convert_ports(cls, v) -> List[PortMapping]: + ports = [] + for i in v: + if isinstance(i, int): + ports.append(PortMapping(container_port=i)) + elif isinstance(i, str): + ports.append(PortMapping.parse(i)) + else: + ports.append(i) + return ports + @validator("env") def convert_env(cls, v) -> Dict[str, str]: if isinstance(v, list): diff --git a/docs/docs/reference/cli/run.md b/docs/docs/reference/cli/run.md index 9edb47f25..36ca1b149 100644 --- a/docs/docs/reference/cli/run.md +++ b/docs/docs/reference/cli/run.md @@ -41,8 +41,8 @@ The following arguments are optional: - `-d`, `--detach` – (Optional) Run in the detached mode. Means, the command doesn't poll logs and run status. -[//]: # (- `-p PORT [PORT ...]`, `--port PORT [PORT ...]` – (Optional) Requests ports or define mappings for them (`APP_PORT:LOCAL_PORT`)) [//]: # (- `-t TAG`, `--tag TAG` – (Optional) A tag name. Warning, if the tag exists, it will be overridden.) +- `-p PORT [PORT ...]`, `--ports PORT [PORT ...]` – (Optional) Requests ports or define mappings for them (`LOCAL_PORT:CONTAINER_PORT`) - `ARGS` – (Optional) Use `ARGS` to pass custom run arguments Spot policy (the arguments are mutually exclusive):