Allow neuro port-forward to accept multiple local-remote port pairs (#…

…695)

CHANGELOG.D/632.feature

@@ -0,0 +1 @@
+`neuro port-forward` command now accepts multiple local-remote port pairs in order to forward several ports by a single command.

README.md

@@ -106,7 +106,7 @@ Name | Description|
 | _[neuro ps](#neuro-ps)_| List all jobs |
 | _[neuro status](#neuro-status)_| Display status of a job |
 | _[neuro exec](#neuro-exec)_| Execute command in a running job |
-| _[neuro port-forward](#neuro-port-forward)_| Forward a port of a running job to a local port |
+| _[neuro port-forward](#neuro-port-forward)_| Forward port\(s) of a running job to local port\(s) |
 | _[neuro logs](#neuro-logs)_| Print the logs for a container |
 | _[neuro kill](#neuro-kill)_| Kill job\(s) |
 | _[neuro top](#neuro-top)_| Display GPU/CPU/Memory usage |
@@ -150,7 +150,7 @@ Name | Description|
 | _[neuro job ls](#neuro-job-ls)_| List all jobs |
 | _[neuro job status](#neuro-job-status)_| Display status of a job |
 | _[neuro job exec](#neuro-job-exec)_| Execute command in a running job |
-| _[neuro job port-forward](#neuro-job-port-forward)_| Forward a port of a running job to a local port |
+| _[neuro job port-forward](#neuro-job-port-forward)_| Forward port\(s) of a running job to local port\(s) |
 | _[neuro job logs](#neuro-job-logs)_| Print the logs for a container |
 | _[neuro job kill](#neuro-job-kill)_| Kill job\(s) |
 | _[neuro job top](#neuro-job-top)_| Display GPU/CPU/Memory usage |
@@ -279,12 +279,12 @@ Name | Description|
 
 ### neuro job port-forward
 
-Forward a port of a running job to a local port.
+Forward port\(s) of a running job to local port\(s).
 
 **Usage:**
 
 ```bash
-neuro job port-forward [OPTIONS] JOB LOCAL_PORT REMOTE_PORT
+neuro job port-forward [OPTIONS] JOB [LOCAL_REMOTE_PORT]...
 ```
 
 **Options:**
@@ -980,12 +980,12 @@ Name | Description|
 
 ## neuro port-forward
 
-Forward a port of a running job to a local port.
+Forward port\(s) of a running job to local port\(s).
 
 **Usage:**
 
 ```bash
-neuro port-forward [OPTIONS] JOB LOCAL_PORT REMOTE_PORT
+neuro port-forward [OPTIONS] JOB [LOCAL_REMOTE_PORT]...
 ```
 
 **Options:**

neuromation/api/jobs.py

@@ -512,11 +512,12 @@ async def port_forward(
                 "UserKnownHostsFile=/dev/null",
             ]
         command += [f"{server_url.user}@{server_url.host}"]
-        print(f"Port of {id} is forwarded to localhost:{local_port}")
-        print(f"Press ^C to stop forwarding")
         proc = await asyncio.create_subprocess_exec(*command)
         try:
-            return await proc.wait()
+            result = await proc.wait()
+            if result != 0:
+                raise ValueError(f"error code {result}")
+            return local_port
         finally:
             await kill_proc_tree(proc.pid, timeout=10)
             # add a sleep to get process watcher a chance to execute all callbacks

neuromation/cli/job.py

@@ -3,7 +3,7 @@
 import os
 import shlex
 import sys
-from typing import Sequence
+from typing import List, Sequence, Tuple
 
 import click
 
@@ -16,6 +16,7 @@
     Resources,
     Volume,
 )
+from neuromation.cli.utils import LOCAL_REMOTE_PORT
 from neuromation.strings.parse import to_megabytes_str
 
 from .defaults import (
@@ -286,26 +287,45 @@ async def exec(
 
 @command(context_settings=dict(ignore_unknown_options=True))
 @click.argument("job")
-@click.argument("local_port", type=int)
-@click.argument("remote_port", type=int)
+@click.argument("local_remote_port", type=LOCAL_REMOTE_PORT, nargs=-1)
 @click.option(
     "--no-key-check",
     is_flag=True,
     help="Disable host key checks. Should be used with caution.",
 )
 @async_cmd
 async def port_forward(
-    cfg: Config, job: str, no_key_check: bool, local_port: int, remote_port: int
+    cfg: Config, job: str, no_key_check: bool, local_remote_port: List[Tuple[int, int]]
 ) -> None:
     """
-    Forward a port of a running job to a local port.
+    Forward port(s) of a running job to local port(s).
     """
+    loop = asyncio.get_event_loop()
     async with cfg.make_client() as client:
-        id = await resolve_job(client, job)
-        retcode = await client.jobs.port_forward(
-            id, no_key_check, local_port, remote_port
-        )
-    sys.exit(retcode)
+        job_id = await resolve_job(client, job)
+        tasks = []
+        for local_port, remote_port in local_remote_port:
+            print(f"Port of {job_id} will be forwarded to localhost:{local_port}")
+            tasks.append(
+                loop.create_task(
+                    client.jobs.port_forward(
+                        job_id, no_key_check, local_port, remote_port
+                    )
+                )
+            )
+
+        print("Press ^C to stop forwarding")
+        result = 0
+        for future in asyncio.as_completed(tasks):
+            try:
+                await future
+            except ValueError as e:
+                print(f"Port forwarding failed: {e}")
+                [task.cancel() for task in tasks]
+                result = -1
+                break
+
+    sys.exit(result)
 
 
 @command()

neuromation/cli/utils.py

@@ -20,6 +20,7 @@
 )
 
 import click
+from click import BadParameter
 from yarl import URL
 
 from neuromation.api import (
@@ -389,3 +390,20 @@ def convert(
 
     def __repr__(self) -> str:
         return "Image"
+
+
+class LocalRemotePortParamType(click.ParamType):
+    def convert(
+        self, value: str, param: Optional[click.Parameter], ctx: Optional[click.Context]
+    ) -> Tuple[int, int]:
+        try:
+            local_str, remote_str = value.split(":")
+            local, remote = int(local_str), int(remote_str)
+            if not (0 < local <= 65535 and 0 < remote <= 65535):
+                raise ValueError("Port should be in range 1 to 65535")
+            return local, remote
+        except ValueError as e:
+            raise BadParameter(f"{value} is not a valid port combination: {e}")
+
+
+LOCAL_REMOTE_PORT = LocalRemotePortParamType()

tests/cli/test_utils.py

@@ -1,9 +1,11 @@
+import click
 import pytest
 from aiohttp import web
 from yarl import URL
 
 from neuromation.api import Action
 from neuromation.cli.utils import (
+    LocalRemotePortParamType,
     parse_permission_action,
     parse_resource_for_sharing,
     resolve_job,
@@ -216,3 +218,32 @@ def test_parse_permission_action_wrong_empty(config):
     err = "invalid permission action '', allowed values: read, write, manage"
     with pytest.raises(ValueError, match=err):
         parse_permission_action(action)
+
+
+@pytest.mark.parametrize(
+    "arg,val",
+    [("1:1", (1, 1)), ("1:10", (1, 10)), ("434:1", (434, 1)), ("0897:123", (897, 123))],
+)
+def test_local_remote_port_param_type_valid(arg, val) -> None:
+    param = LocalRemotePortParamType()
+    assert param.convert(arg, None, None) == val
+
+
+@pytest.mark.parametrize(
+    "arg",
+    [
+        "1:",
+        "-123:10",
+        "34:-65500",
+        "hello:45",
+        "5555:world",
+        "65536:1",
+        "0:0",
+        "none",
+        "",
+    ],
+)
+def test_local_remote_port_param_type_invalid(arg) -> None:
+    param = LocalRemotePortParamType()
+    with pytest.raises(click.BadParameter, match=".* is not a valid port combination"):
+        param.convert(arg, None, None)

tests/e2e/test_e2e_jobs.py

@@ -919,8 +919,15 @@ async def get_(url):
 def test_port_forward_no_job(helper, nginx_job):
     job_name = f"non-existing-job-{uuid4()}"
     with pytest.raises(SystemExit) as cm:
-        helper.run_cli(["port-forward", "--no-key-check", job_name, "0", "0"])
-    assert cm.value.code == 127
+        helper.run_cli(["port-forward", "--no-key-check", job_name, "1:1"])
+    assert cm.value.code == -1
+
+
+@pytest.mark.e2e
+def test_port_forward_invalid_port(helper, nginx_job):
+    with pytest.raises(SystemExit) as cm:
+        helper.run_cli(["port-forward", "--no-key-check", nginx_job, "1:1"])
+    assert cm.value.code == -1
 
 
 @pytest.mark.e2e