Add --ports to dstack run (#573)

* Add --ports to CLI, align port forwarding syntax with docker

* Update --ports in docs

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)

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

cli/dstack/_internal/core/configuration.py

@@ -1,10 +1,12 @@
+import re
 from enum import Enum
 from typing import Dict, List, Optional, Union
 
 from pydantic import BaseModel, Extra, Field, conint, constr, validator
 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):

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):