Tools to connect to and interact with the Mila cluster

Last update: Dec 1, 2022

Related tags

Overview

milatools

The milatools package provides the mila command, which is meant to help with connecting to and interacting with the Mila cluster.

Install

Requires Python >= 3.8

pip install git+git://github.com/mila-iqia/milatools.git

After installing milatools, start with mila init:

mila init

Commands

mila init

Set up your access to the mila cluster interactively. Have your username and password ready!

Set up your SSH config for easy connection with ssh mila
Set up your public key if you don't already have them
Copy your public key over to the cluster for passwordless auth
Set up a public key on the login node to enable ssh into compute nodes

mila code

Connect a VSCode instance to a compute node. mila code first allocates a compute node using slurm (you can pass slurm options as well), and then calls the code command with the appropriate options to start a remote coding session on the allocated node.

You can simply Ctrl+C the process to end the session.

usage: mila code [-h] PATH ...

positional arguments:
  PATH
  SLURM_OPTS

optional arguments:
  -h, --help  show this help message and exit

For example:

mila code path/to/my/experiment

Comments

problem regarding 'mila init'

$ mila init Checking ssh config Traceback (most recent call last): File "/home/pharry/ENV/bin/mila", line 8, in sys.exit(main()) File "/home/pharry/ENV/lib/python3.8/site-packages/milatools/commands.py", line 13, in main auto_cli(milatools) File "/home/pharry/ENV/lib/python3.8/site-packages/coleo/cli.py", line 611, in auto_cli result = run_cli(entry, args, **kwargs) File "/home/pharry/ENV/lib/python3.8/site-packages/coleo/cli.py", line 583, in run_cli return call(opts=opts, args=args) File "/home/pharry/ENV/lib/python3.8/site-packages/coleo/cli.py", line 542, in thunk result = fn(*args) File "/home/pharry/ENV/lib/python3.8/site-packages/ptera/core.py", line 853, in call rval = super().call(*self.partial_args, *args, **kwargs) File "/home/pharry/ENV/lib/python3.8/site-packages/ptera/selfless.py", line 656, in call return self.fn(self, *args, **kwargs) File "/home/pharry/ENV/lib/python3.8/site-packages/milatools/commands.py", line 41, in init c = SSHConfig() File "/home/pharry/ENV/lib/python3.8/site-packages/milatools/utils.py", line 137, in init self.cfg = read_ssh_config(os.path.expanduser("~/.ssh/config")) File "/home/pharry/ENV/lib/python3.8/site-packages/sshconf.py", line 367, in read_ssh_config master_config = read_ssh_config_file(master_path) File "/home/pharry/ENV/lib/python3.8/site-packages/sshconf.py", line 121, in read_ssh_config_file with open(path, "r") as fh_: FileNotFoundError: [Errno 2] No such file or directory: '/home/pharry/.ssh/config'

opened by PwnerHarry 12
Feature request: support VSCode Insiders edition
I'd like to open a ssh connection to a compute node in VSCode Insiders, rather than VSCode (that is, I think as simple as using cli command code-insiders rather than code). From what I can tell, there isn't currently a way to do this, as running

mila code $PATH_ON_CLUSTER

will salloc a node $NODE_NAME, and try to connect with it using the command code -nw --remote ssh-remote+$NODE_NAME $PATH_ON_CLUSTER. I'd like to have the option to use code-insiders ... instead.

To enable this is I think as simple as modifying "code", -> "code-insiders", in milatools/cli/__main__.py (line 277) https://github.com/mila-iqia/milatools/blob/b4c52579391f7674d88a8039eddc1bf8a7cca8f1/milatools/cli/main.py#L276-L280

But ideally, I'd be able to do this with a different command, like

mila code-insiders $PATH_ON_CLUSTER

Is there a reason why this isn't a good idea?
opened by postylem 7

Could not establish connection to compute node (Could not resolve hostname error)

Hello, I am using the mila code command to establish a remote ssh connection with VSCode. It has worked without any problems in the past 3 weeks, but yesterday I started getting this error when running it.

The command I run from my terminal is: mila code <path_to_my_cluster_code> --alloc --cpus-per-task 4 --gres gpu:1 --time=0-3:00:00 --partition unkillable

The issue likely comes from mila code and not from the cluster or ssh because it works fine when I ssh directly without using mila code. Also, I tried setting up the ssh host manually on VSCode (by first getting an salloc on the cluster, then opening VSCode and manually adding the compute node I was given as new host (e.g. ssh -J mila <username>@<nodename>)) and this works perfectly.

The connection fails with the following log on VSCode:

[[17:05:45.046] Log Level: 2
[17:05:45.046] [email protected]
[17:05:45.047] darwin arm64
[17:05:45.055] SSH Resolver called for "ssh-remote+cn-c017.server.mila.quebec", attempt 1
[17:05:45.055] "remote.SSH.useLocalServer": true
[17:05:45.055] "remote.SSH.path": undefined
[17:05:45.055] "remote.SSH.configFile": undefined
[17:05:45.056] "remote.SSH.useFlock": true
[17:05:45.056] "remote.SSH.lockfilesInTmp": false
[17:05:45.056] "remote.SSH.localServerDownload": auto
[17:05:45.056] "remote.SSH.remoteServerListenOnSocket": false
[17:05:45.056] "remote.SSH.showLoginTerminal": false
[17:05:45.056] "remote.SSH.defaultExtensions": []
[17:05:45.056] "remote.SSH.loglevel": 2
[17:05:45.056] "remote.SSH.enableDynamicForwarding": true
[17:05:45.056] "remote.SSH.enableRemoteCommand": false
[17:05:45.056] "remote.SSH.serverPickPortsFromRange": {}
[17:05:45.056] "remote.SSH.serverInstallPath": {}
[17:05:45.058] SSH Resolver called for host: cn-c017.server.mila.quebec
[17:05:45.058] Setting up SSH remote "cn-c017.server.mila.quebec"
[17:05:45.060] Acquiring local install lock: /var/folders/ml/kpt1qf1s2ngfsbzl5ycgclz00000gp/T/vscode-remote-ssh-feed4297-install.lock
[17:05:45.060] Looking for existing server data file at /Users/aldozaimi/Library/Application Support/Code/User/globalStorage/ms-vscode-remote.remote-ssh/vscode-ssh-host-feed4297-3b889b090b5ad5793f524b5d1d39fda662b96a2a-0.84.0/data.json
[17:05:45.060] Using commit id "3b889b090b5ad5793f524b5d1d39fda662b96a2a" and quality "stable" for server
[17:05:45.062] Install and start server if needed
[17:05:45.063] PATH: /opt/homebrew/Caskroom/miniforge/base/envs/simulation_venv/bin:/opt/homebrew/Caskroom/miniforge/base/condabin:/opt/homebrew/bin:/opt/homebrew/sbin:/Library/Frameworks/Python.framework/Versions/3.10/bin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin:/Users/aldozaimi/.cargo/bin
[17:05:45.063] Checking ssh with "ssh -V"
[17:05:45.066] > OpenSSH_8.6p1, LibreSSL 3.3.5

[17:05:45.067] askpass server listening on /var/folders/ml/kpt1qf1s2ngfsbzl5ycgclz00000gp/T/vscode-ssh-askpass-135acdbd1f53ce3ecf04f0db01904fbfecec5929.sock
[17:05:45.068] Spawning local server with {"serverId":1,"ipcHandlePath":"/var/folders/ml/kpt1qf1s2ngfsbzl5ycgclz00000gp/T/vscode-ssh-askpass-d3660326b270bf9a6b665c435ec56bc6679ee114.sock","sshCommand":"ssh","sshArgs":["-v","-T","-D","51100","-o","ConnectTimeout=15","cn-c017.server.mila.quebec"],"serverDataFolderName":".vscode-server","dataFilePath":"/Users/aldozaimi/Library/Application Support/Code/User/globalStorage/ms-vscode-remote.remote-ssh/vscode-ssh-host-feed4297-3b889b090b5ad5793f524b5d1d39fda662b96a2a-0.84.0/data.json"}
[17:05:45.068] Local server env: {"SSH_AUTH_SOCK":"/private/tmp/com.apple.launchd.4M5703FS2F/Listeners","SHELL":"/bin/zsh","DISPLAY":"1","ELECTRON_RUN_AS_NODE":"1","SSH_ASKPASS":"/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/local-server/askpass.sh","VSCODE_SSH_ASKPASS_NODE":"/Applications/Visual Studio Code.app/Contents/Frameworks/Code Helper.app/Contents/MacOS/Code Helper","VSCODE_SSH_ASKPASS_EXTRA_ARGS":"--ms-enable-electron-run-as-node","VSCODE_SSH_ASKPASS_MAIN":"/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/askpass-main.js","VSCODE_SSH_ASKPASS_HANDLE":"/var/folders/ml/kpt1qf1s2ngfsbzl5ycgclz00000gp/T/vscode-ssh-askpass-135acdbd1f53ce3ecf04f0db01904fbfecec5929.sock"}
[17:05:45.068] Spawned 65717
[17:05:45.134] > local-server-1> Spawned ssh, pid=65724
[17:05:45.136] stderr> OpenSSH_8.6p1, LibreSSL 3.3.5
[17:05:45.141] stderr> OpenSSH_8.6p1, LibreSSL 3.3.5
[17:05:45.154] stderr> kex_exchange_identification: Connection closed by remote host
[17:05:45.154] stderr> Connection closed by 172.16.2.25 port 2222
[17:05:45.154] stderr> kex_exchange_identification: Connection closed by remote host
[17:05:45.154] stderr> Connection closed by UNKNOWN port 65535
[17:05:45.155] > local-server-1> ssh child died, shutting down
[17:05:45.156] Local server exit: 0
[17:05:45.156] Received install output: local-server-1> Spawned ssh, pid=65724
OpenSSH_8.6p1, LibreSSL 3.3.5
OpenSSH_8.6p1, LibreSSL 3.3.5
kex_exchange_identification: Connection closed by remote host
Connection closed by 172.16.2.25 port 2222
kex_exchange_identification: Connection closed by remote host
Connection closed by UNKNOWN port 65535
local-server-1> ssh child died, shutting down

[17:05:45.157] Failed to parse remote port from server output
[17:05:45.157] Resolver error: Error: 
	at Function.Create (/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:585222)
	at Object.t.handleInstallOutput (/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:583874)
	at Object.e [as tryInstallWithLocalServer] (/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:624373)
	at processTicksAndRejections (node:internal/process/task_queues:96:5)
	at async /Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:643506
	at async Object.t.withShowDetailsEvent (/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:647224)
	at async /Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:622845
	at async T (/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:619351)
	at async Object.t.resolveWithLocalServer (/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:622460)
	at async Object.t.resolve (/Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:644834)
	at async /Users/aldozaimi/.vscode/extensions/ms-vscode-remote.remote-ssh-0.84.0/out/extension.js:1:727082
[17:05:45.160] ------

Any ideas on how to resolve this?

opened by alzaia 3

command not found: `salloc`

$ mila code ~/ccai/github --alloc --gres=gpu:1 --mem=16G --partition=main
(local) $ ssh mila -fNMS /Users/victor/.ssh/sockets/milatools.mila
(mila) $ salloc --gres=gpu:1 --mem=16G --partition=main
# zsh:1: command not found: salloc
ERROR: Could not find the node name for the allocation

I think a way to solve this is by adding . /etc/profile at login before calling salloc but I don't really know why or what it does, I remember it as a quick fix

opened by vict0rsch 3

mila code