diff --git a/CHANGELOG.D/632.feature b/CHANGELOG.D/632.feature new file mode 100644 index 000000000..33b210e86 --- /dev/null +++ b/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. diff --git a/README.md b/README.md index a013790f7..bf88fedba 100644 --- a/README.md +++ b/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:** diff --git a/neuromation/api/jobs.py b/neuromation/api/jobs.py index 96d8111f1..a5cc952cc 100644 --- a/neuromation/api/jobs.py +++ b/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 diff --git a/neuromation/cli/job.py b/neuromation/cli/job.py index 53dbe2c3e..ad4ef6fec 100644 --- a/neuromation/cli/job.py +++ b/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,8 +287,7 @@ 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, @@ -295,17 +295,37 @@ async def exec( ) @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() diff --git a/neuromation/cli/utils.py b/neuromation/cli/utils.py index 3f82211e1..5c08925b6 100644 --- a/neuromation/cli/utils.py +++ b/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() diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py index a45e4cd4a..0a8ae354b 100644 --- a/tests/cli/test_utils.py +++ b/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) diff --git a/tests/e2e/test_e2e_jobs.py b/tests/e2e/test_e2e_jobs.py index f949ce71a..fe1d81f0d 100644 --- a/tests/e2e/test_e2e_jobs.py +++ b/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