From a5ecdcc25300d9c31e482091e32350a771e0a975 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Mon, 6 Nov 2023 15:51:22 -0500 Subject: [PATCH 01/26] Dont add ControlMaster ssh entries on Windows Signed-off-by: Fabrice Normandin --- milatools/cli/init_command.py | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index 56905c4d..c0ed9fee 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -3,6 +3,7 @@ import difflib from logging import getLogger as get_logger from pathlib import Path +import sys import questionary as qn @@ -39,21 +40,29 @@ def setup_ssh_config( # sure that the directory actually exists. control_path_dir.expanduser().mkdir(exist_ok=True, parents=True) + if sys.platform == "win32": + ssh_multiplexing_config = {} + else: + ssh_multiplexing_config = { + # Tries to reuse an existing connection, but if it fails, it will create a new one. + "ControlMaster": "auto", + # This makes a file per connection, like normandf@login.server.mila.quebec:2222 + "ControlPath": str(control_path_dir / r"%r@%h:%p"), + # persist for 10 minutes after the last connection ends. + "ControlPersist": 600, + } + _add_ssh_entry( ssh_config, - "mila", + host="mila", + Host=None, HostName="login.server.mila.quebec", User=username, PreferredAuthentications="publickey,keyboard-interactive", Port=2222, ServerAliveInterval=120, ServerAliveCountMax=5, - # Tries to reuse an existing connection, but if it fails, it will create a new one. - ControlMaster="auto", - # This makes a file per connection, like normandf@login.server.mila.quebec:2222 - ControlPath=str(control_path_dir / r"%r@%h:%p"), - # persist for 10 minutes after the last connection ends. - ControlPersist=600, + **ssh_multiplexing_config, ) _add_ssh_entry( @@ -97,12 +106,7 @@ def setup_ssh_config( HostName="%h", User=username, ProxyJump="mila", - # Tries to reuse an existing connection, but if it fails, it will create a new one. - ControlMaster="auto", - # This makes a file per connection, like normandf@login.server.mila.quebec:2222 - ControlPath=str(control_path_dir / r"%r@%h:%p"), - # persist for 10 minutes after the last connection ends. - ControlPersist=600, + **ssh_multiplexing_config, ) new_config = ssh_config.cfg.config() @@ -206,6 +210,7 @@ def _add_ssh_entry( ssh_config: SSHConfig, host: str, Host: str | None = None, + *, _space_before: bool = True, _space_after: bool = False, **entry, From b5f6ce9c61dbeea17a3c663f29d18a79dfcb7502 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Mon, 6 Nov 2023 16:44:32 -0500 Subject: [PATCH 02/26] Remove the unused params of _add_ssh_entry Signed-off-by: Fabrice Normandin --- milatools/cli/init_command.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index c0ed9fee..4c783e0d 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -55,7 +55,6 @@ def setup_ssh_config( _add_ssh_entry( ssh_config, host="mila", - Host=None, HostName="login.server.mila.quebec", User=username, PreferredAuthentications="publickey,keyboard-interactive", @@ -210,9 +209,6 @@ def _add_ssh_entry( ssh_config: SSHConfig, host: str, Host: str | None = None, - *, - _space_before: bool = True, - _space_after: bool = False, **entry, ) -> None: """Interactively add an entry to the ssh config file. From e43a4fb77e9bd5e2ff8bda1fa29f0836fd8580ec Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Mon, 6 Nov 2023 16:47:45 -0500 Subject: [PATCH 03/26] Rebase on Master Signed-off-by: Fabrice Normandin --- milatools/cli/init_command.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index 4c783e0d..b4dce40a 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -209,6 +209,9 @@ def _add_ssh_entry( ssh_config: SSHConfig, host: str, Host: str | None = None, + *, + _space_before: bool = True, + _space_after: bool = False, **entry, ) -> None: """Interactively add an entry to the ssh config file. From 05b2b9a09297e856519df3d17ea60e265dc3735c Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Mon, 6 Nov 2023 16:49:13 -0500 Subject: [PATCH 04/26] Fix isort issues Signed-off-by: Fabrice Normandin --- milatools/cli/init_command.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index b4dce40a..b2338cb7 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -1,9 +1,9 @@ from __future__ import annotations import difflib +import sys from logging import getLogger as get_logger from pathlib import Path -import sys import questionary as qn From c6978afabede7c370fd45eb514a69562b67bd8b0 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 2 Nov 2023 21:40:53 +0000 Subject: [PATCH 05/26] Fix the `mila init` command on windows Fixes #63 Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 89c70259..b6d2791d 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -408,6 +408,7 @@ def init(): for entry in os.listdir(sshdir) ): if yn("You have no public keys. Generate one?"): + # TODO: need to get the location of the key as an output of this command! here.run("ssh-keygen") else: exit("No public keys.") @@ -415,10 +416,12 @@ def init(): # Check that it is possible to connect using the key if not here.check_passwordless("mila"): - if yn( - "Your public key does not appear be registered on the cluster. Register it?" - ): - here.run("ssh-copy-id", "mila") + if yn("Your public key does not appear be registered on the cluster. Register it?"): + # NOTE: If we're on a Windows machine, we do something different here: + if sys.platform == "win32": + here.run('cat ~/.ssh/id_rsa.pub | ssh mila "cat >> ~/.ssh/authorized_keys"') + else: + here.run("ssh-copy-id", "mila") if not here.check_passwordless("mila"): exit("ssh-copy-id appears to have failed") else: From fb30bb8fbfe1c7b1f2694e6c3755428d8e6b115c Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Mon, 6 Nov 2023 14:35:37 -0500 Subject: [PATCH 06/26] Fix black formatting errors Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index b6d2791d..e82038b3 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -416,10 +416,14 @@ def init(): # Check that it is possible to connect using the key if not here.check_passwordless("mila"): - if yn("Your public key does not appear be registered on the cluster. Register it?"): + if yn( + "Your public key does not appear be registered on the cluster. Register it?" + ): # NOTE: If we're on a Windows machine, we do something different here: if sys.platform == "win32": - here.run('cat ~/.ssh/id_rsa.pub | ssh mila "cat >> ~/.ssh/authorized_keys"') + here.run( + 'cat ~/.ssh/id_rsa.pub | ssh mila "cat >> ~/.ssh/authorized_keys"' + ) else: here.run("ssh-copy-id", "mila") if not here.check_passwordless("mila"): From 126a46bd61433e28502ef4005263e18cee2ff268 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Mon, 6 Nov 2023 17:08:07 -0500 Subject: [PATCH 07/26] Fix the equivalent of ssh-copy-id, now works! Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index e82038b3..a28a332e 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -421,9 +421,12 @@ def init(): ): # NOTE: If we're on a Windows machine, we do something different here: if sys.platform == "win32": - here.run( - 'cat ~/.ssh/id_rsa.pub | ssh mila "cat >> ~/.ssh/authorized_keys"' + command = ( + "powershell.exe type $env:USERPROFILE\\.ssh\\id_rsa.pub | ssh mila " + '"cat >> ~/.ssh/authorized_keys"' ) + print(T.bold_green("(local) $ " + command)) + subprocess.getoutput(command) else: here.run("ssh-copy-id", "mila") if not here.check_passwordless("mila"): From 557c9ac5534aa5229300144e86c56eaed30901e3 Mon Sep 17 00:00:00 2001 From: Fabrice Date: Tue, 7 Nov 2023 13:33:16 -0500 Subject: [PATCH 08/26] Attempt to fix running `mila code` from WSL Signed-off-by: Fabrice --- milatools/cli/commands.py | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index a28a332e..a77eddb9 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -551,6 +551,11 @@ def code( if command is None: command = os.environ.get("MILATOOLS_CODE_COMMAND", "code") + + # Try to detect if this is being run from within the Windows Subsystem for Linux. + # If so, then we run `code` through a powershell.exe command to open code without + # issues. + running_inside_WSL = bool(shutil.which("powershell.exe")) try: check_disk_quota(remote) @@ -576,15 +581,26 @@ def code( command_path = shutil.which(command) if not command_path: raise CommandNotFoundError(command) + try: while True: - here.run( - command_path, - "-nw", - "--remote", - f"ssh-remote+{qualified(node_name)}", - path, - ) + if running_inside_WSL: + here.run( + "powershell.exe", + command_path, + "-nw", + "--remote", + f"ssh-remote+{qualified(node_name)}", + path, + ) + else: + here.run( + command_path, + "-nw", + "--remote", + f"ssh-remote+{qualified(node_name)}", + path, + ) print( "The editor was closed. Reopen it with " " or terminate the process with " From ddd7ad0c65ddd01dc41e54687949d11f522a3289 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Tue, 7 Nov 2023 14:39:56 -0500 Subject: [PATCH 09/26] Fix whitespace errors and run `code` directly Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index a77eddb9..6b0c6e71 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -551,7 +551,7 @@ def code( if command is None: command = os.environ.get("MILATOOLS_CODE_COMMAND", "code") - + # Try to detect if this is being run from within the Windows Subsystem for Linux. # If so, then we run `code` through a powershell.exe command to open code without # issues. @@ -581,13 +581,13 @@ def code( command_path = shutil.which(command) if not command_path: raise CommandNotFoundError(command) - + try: while True: - if running_inside_WSL: + if running_inside_WSL: here.run( "powershell.exe", - command_path, + "code", "-nw", "--remote", f"ssh-remote+{qualified(node_name)}", From 7f0b0c481dac496a003f92a2a8ab5fff2e3df00a Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Tue, 7 Nov 2023 14:56:49 -0500 Subject: [PATCH 10/26] Set the remote.SSH settings in User Settings Json Signed-off-by: Fabrice Normandin --- milatools/cli/code_command.py | 22 ++++++++++++++++++++++ milatools/cli/commands.py | 21 +++++++++++++++++++-- 2 files changed, 41 insertions(+), 2 deletions(-) create mode 100644 milatools/cli/code_command.py diff --git a/milatools/cli/code_command.py b/milatools/cli/code_command.py new file mode 100644 index 00000000..3e37e57f --- /dev/null +++ b/milatools/cli/code_command.py @@ -0,0 +1,22 @@ +import json +from pathlib import Path +from socket import timeout + + +def set_remote_ssh_vscode_settings( + vscode_settings_json_path: Path, + timeout_seconds: int, + fully_qualified_node_name: str, +) -> None: + # TODO: If on Windows, would need to add this to the settings.json file at this path + # C:\Users\\AppData\Roaming\Code\User\settings.json + # settings_j + with open(vscode_settings_json_path) as f: + settings_json = json.load(f) + settings_json.setdefault("remote.SSH.connectTimeout", timeout_seconds) + remote_platform = settings_json.get("remote.SSH.remotePlatform", {}) + remote_platform.setdefault("fully_qualified_node_name", "linux") + settings_json["remote.SSH.remotePlatform"] = remote_platform + + with open(vscode_settings_json_path, "w") as f: + json.dump(settings_json, f, indent=4) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 6b0c6e71..5aa36d9a 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -27,6 +27,8 @@ from invoke import UnexpectedExit from typing_extensions import TypedDict +from milatools.cli.code_command import set_remote_ssh_vscode_settings + from ..version import version as mversion from .init_command import setup_ssh_config from .local import Local @@ -581,6 +583,21 @@ def code( command_path = shutil.which(command) if not command_path: raise CommandNotFoundError(command) + qualified_node_name = qualified(node_name) + + # TODO: Find the path to this file on other platforms. + vscode_settings_json_path = Path.home() / ( + "AppData\\Roaming\\Code\\User\\settings.json" + if (sys.platform == "win32" or running_inside_WSL) + else "~/.config/Code/User/settings.json" + ) + + if vscode_settings_json_path.exists(): + set_remote_ssh_vscode_settings( + vscode_settings_json_path=vscode_settings_json_path, + timeout_seconds=60, + fully_qualified_node_name=qualified_node_name, + ) try: while True: @@ -590,7 +607,7 @@ def code( "code", "-nw", "--remote", - f"ssh-remote+{qualified(node_name)}", + f"ssh-remote+{qualified_node_name}", path, ) else: @@ -598,7 +615,7 @@ def code( command_path, "-nw", "--remote", - f"ssh-remote+{qualified(node_name)}", + f"ssh-remote+{qualified_node_name}", path, ) print( From 440d1e3955817c5116d1a4dbe806506bb850ebfe Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Tue, 7 Nov 2023 14:58:18 -0500 Subject: [PATCH 11/26] Show that VsCode settings are found Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 1 + 1 file changed, 1 insertion(+) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 5aa36d9a..ae48cdfa 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -593,6 +593,7 @@ def code( ) if vscode_settings_json_path.exists(): + print(T.bold(f"Found VSCode settings file at {vscode_settings_json_path}")) set_remote_ssh_vscode_settings( vscode_settings_json_path=vscode_settings_json_path, timeout_seconds=60, From 1838323068051354a43f03f005cc411a19f9fbff Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Wed, 8 Nov 2023 11:59:10 -0500 Subject: [PATCH 12/26] Remove the VsCode settings thing for now Signed-off-by: Fabrice Normandin --- milatools/cli/code_command.py | 22 ---------------------- milatools/cli/commands.py | 17 ----------------- 2 files changed, 39 deletions(-) delete mode 100644 milatools/cli/code_command.py diff --git a/milatools/cli/code_command.py b/milatools/cli/code_command.py deleted file mode 100644 index 3e37e57f..00000000 --- a/milatools/cli/code_command.py +++ /dev/null @@ -1,22 +0,0 @@ -import json -from pathlib import Path -from socket import timeout - - -def set_remote_ssh_vscode_settings( - vscode_settings_json_path: Path, - timeout_seconds: int, - fully_qualified_node_name: str, -) -> None: - # TODO: If on Windows, would need to add this to the settings.json file at this path - # C:\Users\\AppData\Roaming\Code\User\settings.json - # settings_j - with open(vscode_settings_json_path) as f: - settings_json = json.load(f) - settings_json.setdefault("remote.SSH.connectTimeout", timeout_seconds) - remote_platform = settings_json.get("remote.SSH.remotePlatform", {}) - remote_platform.setdefault("fully_qualified_node_name", "linux") - settings_json["remote.SSH.remotePlatform"] = remote_platform - - with open(vscode_settings_json_path, "w") as f: - json.dump(settings_json, f, indent=4) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index ae48cdfa..0fc1867a 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -27,8 +27,6 @@ from invoke import UnexpectedExit from typing_extensions import TypedDict -from milatools.cli.code_command import set_remote_ssh_vscode_settings - from ..version import version as mversion from .init_command import setup_ssh_config from .local import Local @@ -585,21 +583,6 @@ def code( raise CommandNotFoundError(command) qualified_node_name = qualified(node_name) - # TODO: Find the path to this file on other platforms. - vscode_settings_json_path = Path.home() / ( - "AppData\\Roaming\\Code\\User\\settings.json" - if (sys.platform == "win32" or running_inside_WSL) - else "~/.config/Code/User/settings.json" - ) - - if vscode_settings_json_path.exists(): - print(T.bold(f"Found VSCode settings file at {vscode_settings_json_path}")) - set_remote_ssh_vscode_settings( - vscode_settings_json_path=vscode_settings_json_path, - timeout_seconds=60, - fully_qualified_node_name=qualified_node_name, - ) - try: while True: if running_inside_WSL: From 0b92c8fdac06ec256709878b50707d20ccbfe439 Mon Sep 17 00:00:00 2001 From: Fabrice Date: Tue, 7 Nov 2023 16:28:47 -0500 Subject: [PATCH 13/26] Use `here.run` to run the command Signed-off-by: Fabrice --- milatools/cli/commands.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 0fc1867a..772cf85a 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -425,8 +425,7 @@ def init(): "powershell.exe type $env:USERPROFILE\\.ssh\\id_rsa.pub | ssh mila " '"cat >> ~/.ssh/authorized_keys"' ) - print(T.bold_green("(local) $ " + command)) - subprocess.getoutput(command) + here.run(command) else: here.run("ssh-copy-id", "mila") if not here.check_passwordless("mila"): From a0ad53f8fc67693d0183eec05031a526ec86a222 Mon Sep 17 00:00:00 2001 From: Fabrice Date: Tue, 7 Nov 2023 16:29:09 -0500 Subject: [PATCH 14/26] Skip the hostname checking on Windows machines Signed-off-by: Fabrice --- milatools/cli/commands.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 772cf85a..4a7448a1 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -50,8 +50,7 @@ def main(): - on_mila = get_fully_qualified_name().endswith(".server.mila.quebec") - if on_mila: + if sys.platform != "win32" and get_fully_qualified_name().endswith(".server.mila.quebec"): exit( "ERROR: 'mila ...' should be run on your local machine and not on the Mila cluster" ) From e8817053f9953f011203b6f0a6aab9301ce829c8 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Wed, 8 Nov 2023 12:22:02 -0500 Subject: [PATCH 15/26] Fix black formatting issue Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 4a7448a1..a9f81b02 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -50,7 +50,9 @@ def main(): - if sys.platform != "win32" and get_fully_qualified_name().endswith(".server.mila.quebec"): + if sys.platform != "win32" and get_fully_qualified_name().endswith( + ".server.mila.quebec" + ): exit( "ERROR: 'mila ...' should be run on your local machine and not on the Mila cluster" ) From 477997d5745796e41d02994dee83e70399f55998 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Wed, 8 Nov 2023 13:02:34 -0500 Subject: [PATCH 16/26] Fix missing condition in WSL check Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index a9f81b02..7337e761 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -555,7 +555,9 @@ def code( # Try to detect if this is being run from within the Windows Subsystem for Linux. # If so, then we run `code` through a powershell.exe command to open code without # issues. - running_inside_WSL = bool(shutil.which("powershell.exe")) + running_inside_WSL = sys.platform == "linux" and bool( + shutil.which("powershell.exe") + ) try: check_disk_quota(remote) From 9b5fd377c4ab22c12d5cf2db84331536d75ed3ab Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Wed, 8 Nov 2023 14:41:00 -0500 Subject: [PATCH 17/26] Warn WSL users if `mila init` not done on Windows Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 54 ++++++++++++++++++++++++++++++++++----- 1 file changed, 47 insertions(+), 7 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 7337e761..953f2697 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -5,6 +5,7 @@ from __future__ import annotations import argparse +import functools import operator import os import re @@ -15,6 +16,7 @@ import time import traceback import typing +import warnings import webbrowser from argparse import ArgumentParser, _HelpAction from contextlib import ExitStack @@ -35,6 +37,7 @@ from .utils import ( CommandNotFoundError, MilatoolsUserError, + SSHConfig, SSHConnectionError, T, get_fully_qualified_name, @@ -471,6 +474,12 @@ def init(): else: exit("You will not be able to SSH to a compute node") + # TODO: IF we're running on WSL, we could probably actually just copy the + # id_rsa.pub and the config to the Windows paths (taking care to remove the + # ControlMaster-related entries) so that the user doesn't need to install Python on + # the Windows side. + warn_if_using_WSL_and_mila_init_not_done_on_Windows() + ################### # Welcome message # ################### @@ -552,13 +561,6 @@ def code( if command is None: command = os.environ.get("MILATOOLS_CODE_COMMAND", "code") - # Try to detect if this is being run from within the Windows Subsystem for Linux. - # If so, then we run `code` through a powershell.exe command to open code without - # issues. - running_inside_WSL = sys.platform == "linux" and bool( - shutil.which("powershell.exe") - ) - try: check_disk_quota(remote) except MilatoolsUserError: @@ -585,6 +587,12 @@ def code( raise CommandNotFoundError(command) qualified_node_name = qualified(node_name) + # Try to detect if this is being run from within the Windows Subsystem for Linux. + # If so, then we run `code` through a powershell.exe command to open VSCode without + # issues. + running_inside_WSL = _running_inside_WSL() + warn_if_using_WSL_and_mila_init_not_done_on_Windows() + try: while True: if running_inside_WSL: @@ -624,6 +632,38 @@ def code( print(T.bold(f" ssh mila scancel {data['jobid']}")) +@functools.lru_cache() +def _running_inside_WSL() -> bool: + return sys.platform == "linux" and bool(shutil.which("powershell.exe")) + + +def _mila_init_also_done_on_windows() -> bool: + assert _running_inside_WSL() + windows_username = subprocess.getoutput("powershell.exe '$env:UserName'").strip() + windows_ssh_config_file_path = f"/mnt/c/Users/{windows_username}/.ssh/config" + if not os.path.exists(windows_ssh_config_file_path): + return False + ssh_config = SSHConfig(windows_ssh_config_file_path) + configured_hosts = ssh_config.hosts() + if any(host not in configured_hosts for host in ["mila", "mila-cpu"]): + return False + return True + + +def warn_if_using_WSL_and_mila_init_not_done_on_Windows(): + if _running_inside_WSL() and not _mila_init_also_done_on_windows(): + warnings.warn( + T.orange( + "It seems like you are using the Windows Subsystem for Linux, and " + "haven't yet set-up your SSH config file on the Windows side.\n" + "Make sure to also `pip install milatools` and run `mila init` " + "from a powershell window (assuming you also already installed Python " + "on Windows) so that you can use `mila code` from within WSL without " + "errors." + ) + ) + + def connect(identifier: str, port: int | None): """Reconnect to a persistent server.""" From cbe4d74939c8afe26c964c344571f88d9495f907 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Wed, 8 Nov 2023 16:23:33 -0500 Subject: [PATCH 18/26] Also setup Windows SSH config from WSL Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 50 ++++-------------- milatools/cli/init_command.py | 98 +++++++++++++++++++++++++++++++++-- milatools/cli/utils.py | 9 ++++ 3 files changed, 112 insertions(+), 45 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 953f2697..b8f31d48 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -30,14 +30,16 @@ from typing_extensions import TypedDict from ..version import version as mversion -from .init_command import setup_ssh_config +from .init_command import ( + setup_ssh_config, + setup_windows_ssh_config_from_wsl, +) from .local import Local from .profile import ensure_program, setup_profile from .remote import Remote, SlurmRemote from .utils import ( CommandNotFoundError, MilatoolsUserError, - SSHConfig, SSHConnectionError, T, get_fully_qualified_name, @@ -45,6 +47,7 @@ randname, with_control_file, yn, + running_inside_WSL, ) logger = get_logger(__name__) @@ -390,7 +393,7 @@ def init(): print("Checking ssh config") - setup_ssh_config() + ssh_config = setup_ssh_config() # TODO: Move the rest of this command to functions in the init_command module, # so they can more easily be tested. @@ -478,7 +481,8 @@ def init(): # id_rsa.pub and the config to the Windows paths (taking care to remove the # ControlMaster-related entries) so that the user doesn't need to install Python on # the Windows side. - warn_if_using_WSL_and_mila_init_not_done_on_Windows() + if running_inside_WSL(): + setup_windows_ssh_config_from_wsl(linux_ssh_config=ssh_config) ################### # Welcome message # @@ -590,12 +594,10 @@ def code( # Try to detect if this is being run from within the Windows Subsystem for Linux. # If so, then we run `code` through a powershell.exe command to open VSCode without # issues. - running_inside_WSL = _running_inside_WSL() - warn_if_using_WSL_and_mila_init_not_done_on_Windows() - + inside_WSL = running_inside_WSL() try: while True: - if running_inside_WSL: + if inside_WSL: here.run( "powershell.exe", "code", @@ -632,38 +634,6 @@ def code( print(T.bold(f" ssh mila scancel {data['jobid']}")) -@functools.lru_cache() -def _running_inside_WSL() -> bool: - return sys.platform == "linux" and bool(shutil.which("powershell.exe")) - - -def _mila_init_also_done_on_windows() -> bool: - assert _running_inside_WSL() - windows_username = subprocess.getoutput("powershell.exe '$env:UserName'").strip() - windows_ssh_config_file_path = f"/mnt/c/Users/{windows_username}/.ssh/config" - if not os.path.exists(windows_ssh_config_file_path): - return False - ssh_config = SSHConfig(windows_ssh_config_file_path) - configured_hosts = ssh_config.hosts() - if any(host not in configured_hosts for host in ["mila", "mila-cpu"]): - return False - return True - - -def warn_if_using_WSL_and_mila_init_not_done_on_Windows(): - if _running_inside_WSL() and not _mila_init_also_done_on_windows(): - warnings.warn( - T.orange( - "It seems like you are using the Windows Subsystem for Linux, and " - "haven't yet set-up your SSH config file on the Windows side.\n" - "Make sure to also `pip install milatools` and run `mila init` " - "from a powershell window (assuming you also already installed Python " - "on Windows) so that you can use `mila code` from within WSL without " - "errors." - ) - ) - - def connect(identifier: str, port: int | None): """Reconnect to a persistent server.""" diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index b2338cb7..5468bf36 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -1,20 +1,27 @@ from __future__ import annotations import difflib +import os +import shutil +import subprocess import sys from logging import getLogger as get_logger from pathlib import Path +from typing import Any +import warnings import questionary as qn -from .utils import SSHConfig, T, yn +from .utils import SSHConfig, T, running_inside_WSL, yn + +WINDOWS_UNSUPPORTED_KEYS = ["ControlMaster", "ControlPath", "ControlPersist"] logger = get_logger(__name__) def setup_ssh_config( ssh_config_path: str | Path = "~/.ssh/config", -): +) -> SSHConfig: """Interactively sets up some useful entries in the ~/.ssh/config file on the local machine. Exits if the User cancels any of the prompts or doesn't confirm the changes when asked. @@ -28,6 +35,9 @@ def setup_ssh_config( directly to compute nodes. TODO: Also ask if we should add entries for the ComputeCanada/DRAC clusters. + + Returns: + The resulting SSHConfig if the changes are approved. """ ssh_config_path = _setup_ssh_config_file(ssh_config_path) @@ -116,6 +126,45 @@ def setup_ssh_config( else: ssh_config.save() print(f"Wrote {ssh_config_path}") + return ssh_config + + +def setup_windows_ssh_config_from_wsl(linux_ssh_config: SSHConfig): + """Setup the Windows SSH configuration and public key from within WSL. + + This copies over the entries from the linux ssh configuration file, except for the + values that aren't supported on Windows (e.g. "ControlMaster"). + + This also copies the public key file from the linux SSH directory over to the + Windows SSH directory if it isn't already present. + + This makes it so the user doesn't need to install Python/Anaconda on the Windows + side in order to use `mila code` from within WSL. + """ + assert running_inside_WSL() + windows_home = get_windows_home_path_in_wsl() + windows_ssh_config_path = windows_home / ".ssh/config" + if not windows_ssh_config_path.exists(): + # SSHConfig needs an existing file. + windows_ssh_config_path.touch(mode=0b110_000_000) + + _copy_valid_ssh_entries_to_windows_ssh_config_file( + linux_ssh_config, windows_ssh_config_path + ) + + # copy the public_key_to_windows_ssh_folder + # TODO: This might be different if the user selects a non-default location during + # `ssh-keygen`. + linux_pubkey_file = Path.home() / ".ssh/id_rsa.pub" + windows_pubkey_file = windows_home / ".ssh/id_rsa.pub" + if linux_pubkey_file.exists() and not windows_pubkey_file.exists(): + shutil.copy2(src=linux_pubkey_file, dst=windows_pubkey_file) + + +def get_windows_home_path_in_wsl() -> Path: + assert running_inside_WSL() + windows_username = subprocess.getoutput("powershell.exe '$env:UserName'").strip() + return Path(f"/mnt/c/Users/{windows_username}") def _setup_ssh_config_file(config_file_path: str | Path) -> Path: @@ -150,7 +199,12 @@ def _setup_ssh_config_file(config_file_path: str | Path) -> Path: def _confirm_changes(ssh_config: SSHConfig, previous: str) -> bool: - print(T.bold("The following modifications will be made to your ~/.ssh/config:\n")) + print( + T.bold( + f"The following modifications will be made to your SSH config file at " + f"{ssh_config.path}:\n" + ) + ) diff_lines = list( difflib.unified_diff( (previous + "\n").splitlines(True), @@ -227,7 +281,7 @@ def _add_ssh_entry( existing_entry = ssh_config.host(host) existing_entry.update(entry) ssh_config.cfg.set(host, **existing_entry) - logger.debug(f"Updated {host} entry in ssh config.") + logger.debug(f"Updated {host} entry in ssh config at path {ssh_config.path}.") else: ssh_config.add( host, @@ -235,4 +289,38 @@ def _add_ssh_entry( _space_after=_space_after, **entry, ) - logger.debug(f"Adding new {host} entry in ssh config.") + logger.debug( + f"Adding new {host} entry in ssh config at path {ssh_config.path}." + ) + + +def _copy_valid_ssh_entries_to_windows_ssh_config_file( + linux_ssh_config: SSHConfig, windows_ssh_config_path: Path +): + windows_ssh_config = SSHConfig(windows_ssh_config_path) + initial_windows_config_contents = windows_ssh_config.cfg.config() + + unsupported_keys_lowercase = set(k.lower() for k in WINDOWS_UNSUPPORTED_KEYS) + + for host in linux_ssh_config.hosts(): + linux_ssh_entry: dict[str, Any] = linux_ssh_config.host(host) + _add_ssh_entry( + windows_ssh_config, + host, + **{ + key: value + for key, value in linux_ssh_entry.items() + if key.lower() not in unsupported_keys_lowercase + }, + ) + + new_config_contents = windows_ssh_config.cfg.config() + if new_config_contents == initial_windows_config_contents: + print(f"Did not change ssh config at path {windows_ssh_config.path}") + return + if not _confirm_changes( + windows_ssh_config, previous=initial_windows_config_contents + ): + exit() + # We made changes and they were accepted. + windows_ssh_config.save() diff --git a/milatools/cli/utils.py b/milatools/cli/utils.py index 4c0bdc65..00da8ff3 100644 --- a/milatools/cli/utils.py +++ b/milatools/cli/utils.py @@ -1,13 +1,16 @@ from __future__ import annotations import contextvars +import functools import itertools import random import shlex +import shutil import socket import subprocess from contextlib import contextmanager from pathlib import Path +import sys import blessed import paramiko @@ -140,6 +143,7 @@ class SSHConfig: """Wrapper around sshconf with some extra niceties.""" def __init__(self, path: str | Path): + self.path = path self.cfg = read_ssh_config(path) # self.add = self.cfg.add self.remove = self.cfg.remove @@ -220,3 +224,8 @@ def get_fully_qualified_name() -> str: except Exception: # Fall back, e.g. on Windows. return socket.getfqdn() + + +@functools.lru_cache() +def running_inside_WSL() -> bool: + return sys.platform == "linux" and bool(shutil.which("powershell.exe")) From 611b2b6b224bd40b18bfb2bd4c894706a1375ece Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Wed, 8 Nov 2023 16:26:11 -0500 Subject: [PATCH 19/26] Fix isort issues Signed-off-by: Fabrice Normandin --- milatools/cli/commands.py | 7 ++----- milatools/cli/init_command.py | 2 +- milatools/cli/utils.py | 2 +- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index b8f31d48..b94aa2d8 100644 --- a/milatools/cli/commands.py +++ b/milatools/cli/commands.py @@ -30,10 +30,7 @@ from typing_extensions import TypedDict from ..version import version as mversion -from .init_command import ( - setup_ssh_config, - setup_windows_ssh_config_from_wsl, -) +from .init_command import setup_ssh_config, setup_windows_ssh_config_from_wsl from .local import Local from .profile import ensure_program, setup_profile from .remote import Remote, SlurmRemote @@ -45,9 +42,9 @@ get_fully_qualified_name, qualified, randname, + running_inside_WSL, with_control_file, yn, - running_inside_WSL, ) logger = get_logger(__name__) diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index 5468bf36..b9dd3076 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -5,10 +5,10 @@ import shutil import subprocess import sys +import warnings from logging import getLogger as get_logger from pathlib import Path from typing import Any -import warnings import questionary as qn diff --git a/milatools/cli/utils.py b/milatools/cli/utils.py index 00da8ff3..3efd2ad1 100644 --- a/milatools/cli/utils.py +++ b/milatools/cli/utils.py @@ -8,9 +8,9 @@ import shutil import socket import subprocess +import sys from contextlib import contextmanager from pathlib import Path -import sys import blessed import paramiko From f68784725efb0a4a97422d14824c3168e1f2c335 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 9 Nov 2023 13:10:57 -0500 Subject: [PATCH 20/26] Reuse existing fn to setup windows SSH file Signed-off-by: Fabrice Normandin --- milatools/cli/init_command.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index b9dd3076..dbd4a812 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -144,9 +144,7 @@ def setup_windows_ssh_config_from_wsl(linux_ssh_config: SSHConfig): assert running_inside_WSL() windows_home = get_windows_home_path_in_wsl() windows_ssh_config_path = windows_home / ".ssh/config" - if not windows_ssh_config_path.exists(): - # SSHConfig needs an existing file. - windows_ssh_config_path.touch(mode=0b110_000_000) + windows_ssh_config_path = _setup_ssh_config_file(windows_ssh_config_path) _copy_valid_ssh_entries_to_windows_ssh_config_file( linux_ssh_config, windows_ssh_config_path From 62bd443b23d3ece8906b5845c6d7406906266e9c Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 9 Nov 2023 14:01:39 -0500 Subject: [PATCH 21/26] Preserve ordering of SSH entries when copying Signed-off-by: Fabrice Normandin --- milatools/cli/init_command.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index dbd4a812..4de2c805 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -18,6 +18,9 @@ logger = get_logger(__name__) +HOSTS = ["mila", "mila-cpu", "*.server.mila.quebec !*login.server.mila.quebec"] +"""List of host entries that get added to the SSH configurtion by `mila init`.""" + def setup_ssh_config( ssh_config_path: str | Path = "~/.ssh/config", @@ -156,6 +159,7 @@ def setup_windows_ssh_config_from_wsl(linux_ssh_config: SSHConfig): linux_pubkey_file = Path.home() / ".ssh/id_rsa.pub" windows_pubkey_file = windows_home / ".ssh/id_rsa.pub" if linux_pubkey_file.exists() and not windows_pubkey_file.exists(): + print(f"Copying public key at {linux_pubkey_file} to Windows ssh folder.") shutil.copy2(src=linux_pubkey_file, dst=windows_pubkey_file) @@ -300,7 +304,17 @@ def _copy_valid_ssh_entries_to_windows_ssh_config_file( unsupported_keys_lowercase = set(k.lower() for k in WINDOWS_UNSUPPORTED_KEYS) - for host in linux_ssh_config.hosts(): + # NOTE: need to preserve the ordering of entries: + for host in HOSTS + [ + host for host in linux_ssh_config.hosts() if host not in HOSTS + ]: + if host not in linux_ssh_config.hosts(): + warnings.warn( + RuntimeWarning( + f"Weird, we expected to have a {host!r} entry in the SSH config..." + ) + ) + continue linux_ssh_entry: dict[str, Any] = linux_ssh_config.host(host) _add_ssh_entry( windows_ssh_config, From a2f4b86fa00c3a07a1376d9993970d3b45432cb4 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 9 Nov 2023 14:20:46 -0500 Subject: [PATCH 22/26] Copy both public and private key files Signed-off-by: Fabrice Normandin --- milatools/cli/init_command.py | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index 4de2c805..47084a2c 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -145,6 +145,8 @@ def setup_windows_ssh_config_from_wsl(linux_ssh_config: SSHConfig): side in order to use `mila code` from within WSL. """ assert running_inside_WSL() + # NOTE: This also assumes that a public/private key pair has already been generated + # at ~/.ssh/id_rsa.pub and ~/.ssh/id_rsa. windows_home = get_windows_home_path_in_wsl() windows_ssh_config_path = windows_home / ".ssh/config" windows_ssh_config_path = _setup_ssh_config_file(windows_ssh_config_path) @@ -156,11 +158,27 @@ def setup_windows_ssh_config_from_wsl(linux_ssh_config: SSHConfig): # copy the public_key_to_windows_ssh_folder # TODO: This might be different if the user selects a non-default location during # `ssh-keygen`. - linux_pubkey_file = Path.home() / ".ssh/id_rsa.pub" - windows_pubkey_file = windows_home / ".ssh/id_rsa.pub" - if linux_pubkey_file.exists() and not windows_pubkey_file.exists(): - print(f"Copying public key at {linux_pubkey_file} to Windows ssh folder.") - shutil.copy2(src=linux_pubkey_file, dst=windows_pubkey_file) + + linux_private_key_file = Path.home() / ".ssh/id_rsa" + windows_private_key_file = windows_home / ".ssh/id_rsa" + + for linux_key_file, windows_key_file in [ + (linux_private_key_file, windows_private_key_file), + ( + linux_private_key_file.with_suffix(".pub"), + windows_private_key_file.with_suffix(".pub"), + ), + ]: + _copy_if_needed(linux_key_file, windows_key_file) + + +def _copy_if_needed(linux_key_file: Path, windows_key_file: Path): + if linux_key_file.exists() and not windows_key_file.exists(): + print( + f"Copying {linux_key_file} over to the Windows ssh folder at " + f"{windows_key_file}." + ) + shutil.copy2(src=linux_key_file, dst=windows_key_file) def get_windows_home_path_in_wsl() -> Path: From 90b85b8cb93c3065d3c4c302ef715580b5486d4c Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 9 Nov 2023 15:04:34 -0500 Subject: [PATCH 23/26] Add tests for the Windows SSH setup from WSL Signed-off-by: Fabrice Normandin --- tests/cli/test_init_command.py | 129 +++++++++++++++++- ...tup_windows_ssh_config_from_wsl_accept_.md | 68 +++++++++ ...tup_windows_ssh_config_from_wsl_reject_.md | 68 +++++++++ 3 files changed, 264 insertions(+), 1 deletion(-) create mode 100644 tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_accept_.md create mode 100644 tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_reject_.md diff --git a/tests/cli/test_init_command.py b/tests/cli/test_init_command.py index 91f309f6..4e71f295 100644 --- a/tests/cli/test_init_command.py +++ b/tests/cli/test_init_command.py @@ -4,18 +4,23 @@ import textwrap from functools import partial from pathlib import Path +from unittest.mock import Mock import pytest import questionary from prompt_toolkit.input import PipeInput, create_pipe_input +from pytest_mock import MockerFixture from pytest_regressions.file_regression import FileRegressionFixture +from milatools.cli import init_command from milatools.cli.init_command import ( _get_username, _setup_ssh_config_file, + get_windows_home_path_in_wsl, setup_ssh_config, + setup_windows_ssh_config_from_wsl, ) -from milatools.cli.utils import SSHConfig +from milatools.cli.utils import SSHConfig, running_inside_WSL @pytest.fixture @@ -499,3 +504,125 @@ def test_fixes_dir_permission_issues( assert file.parent.stat().st_mode & 0o777 == 0o700 assert file.exists() assert file.stat().st_mode & 0o777 == 0o600 + + +@pytest.fixture +def linux_ssh_config( + tmp_path: Path, input_pipe: PipeInput, monkeypatch: pytest.MonkeyPatch +) -> SSHConfig: + """Creates the SSH config that would be generated by `mila init`.""" + # Enter username, accept fixing that entry, then confirm. + ssh_config_path = tmp_path / "ssh_config" + + for prompt in ["y", "bob\r", "y", "y", "y", "y", "y"]: + input_pipe.send_text(prompt) + + monkeypatch.setattr("sys.platform", "linux") + # TODO: The config will be different if we run the tests on Windows, it won't + # contain the ControlMaster entries. + setup_ssh_config(ssh_config_path) + + return SSHConfig(ssh_config_path) + + +@pytest.mark.parametrize("user_inputs", [["y"], ["n"]], ids=["accept", "reject"]) +def test_setup_windows_ssh_config_from_wsl( + mocker: MockerFixture, + tmp_path: Path, + linux_ssh_config: SSHConfig, + input_pipe: PipeInput, + file_regression: FileRegressionFixture, + user_inputs: list[str], +): + initial_contents = linux_ssh_config.cfg.config() + windows_home = tmp_path / "fake_windows_home" + windows_home.mkdir(exist_ok=False) + windows_ssh_config_path = windows_home / ".ssh" / "config" + + init_command.running_inside_WSL = mocker.Mock( + spec=running_inside_WSL, return_value=True + ) + init_command.get_windows_home_path_in_wsl = mocker.Mock( + spec=get_windows_home_path_in_wsl, return_value=windows_home + ) + + for prompt in user_inputs: + input_pipe.send_text(prompt) + + setup_windows_ssh_config_from_wsl(linux_ssh_config=linux_ssh_config) + + assert windows_ssh_config_path.exists() + assert windows_ssh_config_path.stat().st_mode & 0o777 == 0o600 + assert windows_ssh_config_path.parent.stat().st_mode & 0o777 == 0o700 + + expected_text = "\n".join( + [ + "When this SSH config is already present in the WSL environment with " + + ( + "\n".join( + [ + "these initial contents:", + "```", + initial_contents, + "```", + "", + ] + ) + if initial_contents.strip() + else "no initial ssh config file" + ), + "", + f"and these user inputs: {user_inputs}", + "leads the following ssh config file on the Windows side:", + "", + "```", + windows_ssh_config_path.read_text(), + "```", + ] + ) + + file_regression.check(expected_text, extension=".md") + + +def test_setup_windows_ssh_config_from_wsl_copies_keys( + mocker: MockerFixture, + tmp_path: Path, + linux_ssh_config: SSHConfig, + input_pipe: PipeInput, + monkeypatch: pytest.MonkeyPatch, +): + linux_home = tmp_path / "fake_linux_home" + linux_home.mkdir(exist_ok=False) + windows_home = tmp_path / "fake_windows_home" + windows_home.mkdir(exist_ok=False) + monkeypatch.setattr(Path, "home", Mock(spec=Path.home, return_value=linux_home)) + + init_command.running_inside_WSL = mocker.Mock( + spec=running_inside_WSL, return_value=True + ) + init_command.get_windows_home_path_in_wsl = mocker.Mock( + spec=get_windows_home_path_in_wsl, return_value=windows_home + ) + + fake_linux_ssh_dir = linux_home / ".ssh" + fake_linux_ssh_dir.mkdir(mode=0o700) + + private_key_text = "THIS IS A PRIVATE KEY" + linux_private_key_path = fake_linux_ssh_dir / "id_rsa" + linux_private_key_path.write_text(private_key_text) + + public_key_text = "THIS IS A PUBLIC KEY" + linux_public_key_path = linux_private_key_path.with_suffix(".pub") + linux_public_key_path.write_text(public_key_text) + + input_pipe.send_text("y") + + setup_windows_ssh_config_from_wsl(linux_ssh_config=linux_ssh_config) + + windows_private_key_path = windows_home / ".ssh" / "id_rsa" + windows_public_key_path = windows_private_key_path.with_suffix(".pub") + + assert windows_private_key_path.exists() + assert windows_private_key_path.read_text() == private_key_text + assert windows_public_key_path.exists() + assert windows_public_key_path.read_text() == public_key_text diff --git a/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_accept_.md b/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_accept_.md new file mode 100644 index 00000000..747efb21 --- /dev/null +++ b/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_accept_.md @@ -0,0 +1,68 @@ +When this SSH config is already present in the WSL environment with these initial contents: +``` + +Host mila + HostName login.server.mila.quebec + User bob + PreferredAuthentications publickey,keyboard-interactive + Port 2222 + ServerAliveInterval 120 + ServerAliveCountMax 5 + ControlMaster auto + ControlPath ~/.cache/ssh/%r@%h:%p + ControlPersist 600 + +Host mila-cpu + User bob + Port 2222 + ForwardAgent yes + StrictHostKeyChecking no + LogLevel ERROR + UserKnownHostsFile /dev/null + RequestTTY force + ConnectTimeout 600 + ServerAliveInterval 120 + ProxyCommand ssh mila "/cvmfs/config.mila.quebec/scripts/milatools/slurm-proxy.sh mila-cpu --mem=8G" + RemoteCommand /cvmfs/config.mila.quebec/scripts/milatools/entrypoint.sh mila-cpu + +Host *.server.mila.quebec !*login.server.mila.quebec + HostName %h + User bob + ProxyJump mila + ControlMaster auto + ControlPath ~/.cache/ssh/%r@%h:%p + ControlPersist 600 +``` + + +and these user inputs: ['y'] +leads the following ssh config file on the Windows side: + +``` + +Host mila + HostName login.server.mila.quebec + User bob + PreferredAuthentications publickey,keyboard-interactive + Port 2222 + ServerAliveInterval 120 + ServerAliveCountMax 5 + +Host mila-cpu + User bob + Port 2222 + ForwardAgent yes + StrictHostKeyChecking no + LogLevel ERROR + UserKnownHostsFile /dev/null + RequestTTY force + ConnectTimeout 600 + ServerAliveInterval 120 + ProxyCommand ssh mila "/cvmfs/config.mila.quebec/scripts/milatools/slurm-proxy.sh mila-cpu --mem=8G" + remotecommand /cvmfs/config.mila.quebec/scripts/milatools/entrypoint.sh mila-cpu + +Host *.server.mila.quebec !*login.server.mila.quebec + HostName %h + User bob + ProxyJump mila +``` \ No newline at end of file diff --git a/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_reject_.md b/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_reject_.md new file mode 100644 index 00000000..50f01ed9 --- /dev/null +++ b/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_reject_.md @@ -0,0 +1,68 @@ +When this SSH config is already present in the WSL environment with these initial contents: +``` + +Host mila + HostName login.server.mila.quebec + User bob + PreferredAuthentications publickey,keyboard-interactive + Port 2222 + ServerAliveInterval 120 + ServerAliveCountMax 5 + ControlMaster auto + ControlPath ~/.cache/ssh/%r@%h:%p + ControlPersist 600 + +Host mila-cpu + User bob + Port 2222 + ForwardAgent yes + StrictHostKeyChecking no + LogLevel ERROR + UserKnownHostsFile /dev/null + RequestTTY force + ConnectTimeout 600 + ServerAliveInterval 120 + ProxyCommand ssh mila "/cvmfs/config.mila.quebec/scripts/milatools/slurm-proxy.sh mila-cpu --mem=8G" + RemoteCommand /cvmfs/config.mila.quebec/scripts/milatools/entrypoint.sh mila-cpu + +Host *.server.mila.quebec !*login.server.mila.quebec + HostName %h + User bob + ProxyJump mila + ControlMaster auto + ControlPath ~/.cache/ssh/%r@%h:%p + ControlPersist 600 +``` + + +and these user inputs: ['n'] +leads the following ssh config file on the Windows side: + +``` + +Host mila + HostName login.server.mila.quebec + User bob + PreferredAuthentications publickey,keyboard-interactive + Port 2222 + ServerAliveInterval 120 + ServerAliveCountMax 5 + +Host mila-cpu + User bob + Port 2222 + ForwardAgent yes + StrictHostKeyChecking no + LogLevel ERROR + UserKnownHostsFile /dev/null + RequestTTY force + ConnectTimeout 600 + ServerAliveInterval 120 + ProxyCommand ssh mila "/cvmfs/config.mila.quebec/scripts/milatools/slurm-proxy.sh mila-cpu --mem=8G" + remotecommand /cvmfs/config.mila.quebec/scripts/milatools/entrypoint.sh mila-cpu + +Host *.server.mila.quebec !*login.server.mila.quebec + HostName %h + User bob + ProxyJump mila +``` \ No newline at end of file From dd877876c71e8ffa1ef2dfaca2885b81343d23e0 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 9 Nov 2023 15:29:46 -0500 Subject: [PATCH 24/26] Remove need for pytest-mock dependency Signed-off-by: Fabrice Normandin --- tests/cli/test_init_command.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/tests/cli/test_init_command.py b/tests/cli/test_init_command.py index 4e71f295..548bab63 100644 --- a/tests/cli/test_init_command.py +++ b/tests/cli/test_init_command.py @@ -9,7 +9,6 @@ import pytest import questionary from prompt_toolkit.input import PipeInput, create_pipe_input -from pytest_mock import MockerFixture from pytest_regressions.file_regression import FileRegressionFixture from milatools.cli import init_command @@ -527,11 +526,11 @@ def linux_ssh_config( @pytest.mark.parametrize("user_inputs", [["y"], ["n"]], ids=["accept", "reject"]) def test_setup_windows_ssh_config_from_wsl( - mocker: MockerFixture, tmp_path: Path, linux_ssh_config: SSHConfig, input_pipe: PipeInput, file_regression: FileRegressionFixture, + monkeypatch: pytest.MonkeyPatch, user_inputs: list[str], ): initial_contents = linux_ssh_config.cfg.config() @@ -539,13 +538,16 @@ def test_setup_windows_ssh_config_from_wsl( windows_home.mkdir(exist_ok=False) windows_ssh_config_path = windows_home / ".ssh" / "config" - init_command.running_inside_WSL = mocker.Mock( - spec=running_inside_WSL, return_value=True + monkeypatch.setattr( + init_command, + running_inside_WSL.__name__, + Mock(spec=running_inside_WSL, return_value=True), ) - init_command.get_windows_home_path_in_wsl = mocker.Mock( - spec=get_windows_home_path_in_wsl, return_value=windows_home + monkeypatch.setattr( + init_command, + get_windows_home_path_in_wsl.__name__, + Mock(spec=get_windows_home_path_in_wsl, return_value=windows_home), ) - for prompt in user_inputs: input_pipe.send_text(prompt) @@ -585,7 +587,6 @@ def test_setup_windows_ssh_config_from_wsl( def test_setup_windows_ssh_config_from_wsl_copies_keys( - mocker: MockerFixture, tmp_path: Path, linux_ssh_config: SSHConfig, input_pipe: PipeInput, @@ -597,11 +598,15 @@ def test_setup_windows_ssh_config_from_wsl_copies_keys( windows_home.mkdir(exist_ok=False) monkeypatch.setattr(Path, "home", Mock(spec=Path.home, return_value=linux_home)) - init_command.running_inside_WSL = mocker.Mock( - spec=running_inside_WSL, return_value=True + monkeypatch.setattr( + init_command, + running_inside_WSL.__name__, + Mock(spec=running_inside_WSL, return_value=True), ) - init_command.get_windows_home_path_in_wsl = mocker.Mock( - spec=get_windows_home_path_in_wsl, return_value=windows_home + monkeypatch.setattr( + init_command, + get_windows_home_path_in_wsl.__name__, + Mock(spec=get_windows_home_path_in_wsl, return_value=windows_home), ) fake_linux_ssh_dir = linux_home / ".ssh" From 531b89a26cc093eb1cf1b303b59c753762d275be Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 9 Nov 2023 15:58:48 -0500 Subject: [PATCH 25/26] Fix a bug in the added tests Signed-off-by: Fabrice Normandin --- .pre-commit-config.yaml | 93 +++++++++++++++++++ milatools/cli/init_command.py | 39 ++++---- tests/cli/test_init_command.py | 19 +++- ...tup_windows_ssh_config_from_wsl_accept_.md | 2 +- ...tup_windows_ssh_config_from_wsl_reject_.md | 27 +----- 5 files changed, 128 insertions(+), 52 deletions(-) create mode 100644 .pre-commit-config.yaml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..f2bcbe23 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,93 @@ +default_language_version: + python: python3 + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.4.0 + hooks: + # list of supported hooks: https://pre-commit.com/hooks.html + - id: trailing-whitespace + require_serial: true + - id: end-of-file-fixer + require_serial: true + # - id: check-docstring-first + - id: check-yaml + require_serial: true + - id: debug-statements + require_serial: true + - id: detect-private-key + require_serial: true + - id: check-executables-have-shebangs + require_serial: true + - id: check-toml + require_serial: true + - id: check-case-conflict + require_serial: true + - id: check-added-large-files + require_serial: true + + - repo: https://github.com/charliermarsh/ruff-pre-commit + # Ruff version. + rev: 'v0.0.261' + hooks: + - id: ruff + # args: ['--fix'] + args: ['--line-length', '88', '--fix'] + require_serial: true + + - repo: https://github.com/pycqa/isort + rev: 5.12.0 + hooks: + - id: isort + name: isort (python) + + # python code formatting + - repo: https://github.com/psf/black + rev: 22.12.0 + hooks: + - id: black + # args: [--line-length, "99"] + args: [--line-length, "88"] + require_serial: true + + # python docstring formatting + - repo: https://github.com/myint/docformatter + rev: v1.5.1 + hooks: + - id: docformatter + # args: [--in-place, --wrap-summaries=99, --wrap-descriptions=99] + args: [--in-place, --wrap-summaries=88, --wrap-descriptions=88] + require_serial: true + + # NOTE: Disabling this, since I'm having the glib-c2.29 weird bug. + # # yaml formatting + # - repo: https://github.com/pre-commit/mirrors-prettier + # rev: v2.7.1 + # hooks: + # - id: prettier + # types: [yaml] + + # md formatting + - repo: https://github.com/executablebooks/mdformat + rev: 0.7.16 + hooks: + - id: mdformat + args: ["--number"] + additional_dependencies: + - mdformat-gfm + - mdformat-tables + - mdformat_frontmatter + # - mdformat-toc + # - mdformat-black + require_serial: true + + + # word spelling linter + - repo: https://github.com/codespell-project/codespell + rev: v2.2.2 + hooks: + - id: codespell + args: + - --skip=logs/**,data/** + # - --ignore-words-list=abc,def + require_serial: true diff --git a/milatools/cli/init_command.py b/milatools/cli/init_command.py index 47084a2c..3e55acfa 100644 --- a/milatools/cli/init_command.py +++ b/milatools/cli/init_command.py @@ -1,7 +1,6 @@ from __future__ import annotations import difflib -import os import shutil import subprocess import sys @@ -151,14 +150,28 @@ def setup_windows_ssh_config_from_wsl(linux_ssh_config: SSHConfig): windows_ssh_config_path = windows_home / ".ssh/config" windows_ssh_config_path = _setup_ssh_config_file(windows_ssh_config_path) + windows_ssh_config = SSHConfig(windows_ssh_config_path) + + initial_windows_config_contents = windows_ssh_config.cfg.config() _copy_valid_ssh_entries_to_windows_ssh_config_file( - linux_ssh_config, windows_ssh_config_path + linux_ssh_config, windows_ssh_config ) + new_windows_config_contents = windows_ssh_config.cfg.config() - # copy the public_key_to_windows_ssh_folder - # TODO: This might be different if the user selects a non-default location during - # `ssh-keygen`. + if ( + new_windows_config_contents != initial_windows_config_contents + and _confirm_changes(windows_ssh_config, initial_windows_config_contents) + ): + # We made changes and they were accepted. + windows_ssh_config.save() + else: + print(f"Did not change ssh config at path {windows_ssh_config.path}") + return # also skip copying the SSH keys. + # Copy the SSH key to the windows folder so that passwordless SSH also works on + # Windows. + # TODO: This will need to change if we support using a non-default location at some + # point. linux_private_key_file = Path.home() / ".ssh/id_rsa" windows_private_key_file = windows_home / ".ssh/id_rsa" @@ -315,11 +328,8 @@ def _add_ssh_entry( def _copy_valid_ssh_entries_to_windows_ssh_config_file( - linux_ssh_config: SSHConfig, windows_ssh_config_path: Path + linux_ssh_config: SSHConfig, windows_ssh_config: SSHConfig ): - windows_ssh_config = SSHConfig(windows_ssh_config_path) - initial_windows_config_contents = windows_ssh_config.cfg.config() - unsupported_keys_lowercase = set(k.lower() for k in WINDOWS_UNSUPPORTED_KEYS) # NOTE: need to preserve the ordering of entries: @@ -343,14 +353,3 @@ def _copy_valid_ssh_entries_to_windows_ssh_config_file( if key.lower() not in unsupported_keys_lowercase }, ) - - new_config_contents = windows_ssh_config.cfg.config() - if new_config_contents == initial_windows_config_contents: - print(f"Did not change ssh config at path {windows_ssh_config.path}") - return - if not _confirm_changes( - windows_ssh_config, previous=initial_windows_config_contents - ): - exit() - # We made changes and they were accepted. - windows_ssh_config.save() diff --git a/tests/cli/test_init_command.py b/tests/cli/test_init_command.py index 548bab63..1dab0b94 100644 --- a/tests/cli/test_init_command.py +++ b/tests/cli/test_init_command.py @@ -513,7 +513,7 @@ def linux_ssh_config( # Enter username, accept fixing that entry, then confirm. ssh_config_path = tmp_path / "ssh_config" - for prompt in ["y", "bob\r", "y", "y", "y", "y", "y"]: + for prompt in ["y", "bob\r", "y"]: input_pipe.send_text(prompt) monkeypatch.setattr("sys.platform", "linux") @@ -524,14 +524,14 @@ def linux_ssh_config( return SSHConfig(ssh_config_path) -@pytest.mark.parametrize("user_inputs", [["y"], ["n"]], ids=["accept", "reject"]) +@pytest.mark.parametrize("accept_changes", [True, False], ids=["accept", "reject"]) def test_setup_windows_ssh_config_from_wsl( tmp_path: Path, linux_ssh_config: SSHConfig, input_pipe: PipeInput, file_regression: FileRegressionFixture, monkeypatch: pytest.MonkeyPatch, - user_inputs: list[str], + accept_changes: bool, ): initial_contents = linux_ssh_config.cfg.config() windows_home = tmp_path / "fake_windows_home" @@ -548,6 +548,12 @@ def test_setup_windows_ssh_config_from_wsl( get_windows_home_path_in_wsl.__name__, Mock(spec=get_windows_home_path_in_wsl, return_value=windows_home), ) + user_inputs: list[str] = [] + if not windows_ssh_config_path.exists(): + # We accept creating the Windows SSH config file for now. + user_inputs.append("y") + user_inputs.append("y" if accept_changes else "n") + for prompt in user_inputs: input_pipe.send_text(prompt) @@ -556,6 +562,8 @@ def test_setup_windows_ssh_config_from_wsl( assert windows_ssh_config_path.exists() assert windows_ssh_config_path.stat().st_mode & 0o777 == 0o600 assert windows_ssh_config_path.parent.stat().st_mode & 0o777 == 0o700 + if not accept_changes: + assert windows_ssh_config_path.read_text() == "" expected_text = "\n".join( [ @@ -574,7 +582,7 @@ def test_setup_windows_ssh_config_from_wsl( else "no initial ssh config file" ), "", - f"and these user inputs: {user_inputs}", + f"and these user inputs: {tuple(user_inputs)}", "leads the following ssh config file on the Windows side:", "", "```", @@ -620,7 +628,8 @@ def test_setup_windows_ssh_config_from_wsl_copies_keys( linux_public_key_path = linux_private_key_path.with_suffix(".pub") linux_public_key_path.write_text(public_key_text) - input_pipe.send_text("y") + input_pipe.send_text("y") # accept creating the Windows config file + input_pipe.send_text("y") # accept the changes setup_windows_ssh_config_from_wsl(linux_ssh_config=linux_ssh_config) diff --git a/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_accept_.md b/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_accept_.md index 747efb21..730cab85 100644 --- a/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_accept_.md +++ b/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_accept_.md @@ -35,7 +35,7 @@ Host *.server.mila.quebec !*login.server.mila.quebec ``` -and these user inputs: ['y'] +and these user inputs: ('y', 'y') leads the following ssh config file on the Windows side: ``` diff --git a/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_reject_.md b/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_reject_.md index 50f01ed9..5940aa01 100644 --- a/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_reject_.md +++ b/tests/cli/test_init_command/test_setup_windows_ssh_config_from_wsl_reject_.md @@ -35,34 +35,9 @@ Host *.server.mila.quebec !*login.server.mila.quebec ``` -and these user inputs: ['n'] +and these user inputs: ('y', 'n') leads the following ssh config file on the Windows side: ``` -Host mila - HostName login.server.mila.quebec - User bob - PreferredAuthentications publickey,keyboard-interactive - Port 2222 - ServerAliveInterval 120 - ServerAliveCountMax 5 - -Host mila-cpu - User bob - Port 2222 - ForwardAgent yes - StrictHostKeyChecking no - LogLevel ERROR - UserKnownHostsFile /dev/null - RequestTTY force - ConnectTimeout 600 - ServerAliveInterval 120 - ProxyCommand ssh mila "/cvmfs/config.mila.quebec/scripts/milatools/slurm-proxy.sh mila-cpu --mem=8G" - remotecommand /cvmfs/config.mila.quebec/scripts/milatools/entrypoint.sh mila-cpu - -Host *.server.mila.quebec !*login.server.mila.quebec - HostName %h - User bob - ProxyJump mila ``` \ No newline at end of file From f073c9edcf7ff64595548e534fdd1d27b3aa846d Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 9 Nov 2023 16:13:12 -0500 Subject: [PATCH 26/26] Remove accidentally-added .pre-commit-config.yaml Signed-off-by: Fabrice Normandin --- .pre-commit-config.yaml | 93 ----------------------------------------- 1 file changed, 93 deletions(-) delete mode 100644 .pre-commit-config.yaml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml deleted file mode 100644 index f2bcbe23..00000000 --- a/.pre-commit-config.yaml +++ /dev/null @@ -1,93 +0,0 @@ -default_language_version: - python: python3 - -repos: - - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.4.0 - hooks: - # list of supported hooks: https://pre-commit.com/hooks.html - - id: trailing-whitespace - require_serial: true - - id: end-of-file-fixer - require_serial: true - # - id: check-docstring-first - - id: check-yaml - require_serial: true - - id: debug-statements - require_serial: true - - id: detect-private-key - require_serial: true - - id: check-executables-have-shebangs - require_serial: true - - id: check-toml - require_serial: true - - id: check-case-conflict - require_serial: true - - id: check-added-large-files - require_serial: true - - - repo: https://github.com/charliermarsh/ruff-pre-commit - # Ruff version. - rev: 'v0.0.261' - hooks: - - id: ruff - # args: ['--fix'] - args: ['--line-length', '88', '--fix'] - require_serial: true - - - repo: https://github.com/pycqa/isort - rev: 5.12.0 - hooks: - - id: isort - name: isort (python) - - # python code formatting - - repo: https://github.com/psf/black - rev: 22.12.0 - hooks: - - id: black - # args: [--line-length, "99"] - args: [--line-length, "88"] - require_serial: true - - # python docstring formatting - - repo: https://github.com/myint/docformatter - rev: v1.5.1 - hooks: - - id: docformatter - # args: [--in-place, --wrap-summaries=99, --wrap-descriptions=99] - args: [--in-place, --wrap-summaries=88, --wrap-descriptions=88] - require_serial: true - - # NOTE: Disabling this, since I'm having the glib-c2.29 weird bug. - # # yaml formatting - # - repo: https://github.com/pre-commit/mirrors-prettier - # rev: v2.7.1 - # hooks: - # - id: prettier - # types: [yaml] - - # md formatting - - repo: https://github.com/executablebooks/mdformat - rev: 0.7.16 - hooks: - - id: mdformat - args: ["--number"] - additional_dependencies: - - mdformat-gfm - - mdformat-tables - - mdformat_frontmatter - # - mdformat-toc - # - mdformat-black - require_serial: true - - - # word spelling linter - - repo: https://github.com/codespell-project/codespell - rev: v2.2.2 - hooks: - - id: codespell - args: - - --skip=logs/**,data/** - # - --ignore-words-list=abc,def - require_serial: true