mila-iqia
diff --git a/‎.github/workflows/build.yml‎
Lines changed: 42 additions & 41 deletions b/‎.github/workflows/build.yml‎
Lines changed: 42 additions & 41 deletions
diff --git a/‎milatools/cli/code.py‎
Lines changed: 26 additions & 1 deletion b/‎milatools/cli/code.py‎
Lines changed: 26 additions & 1 deletion
diff --git a/‎milatools/cli/commands.py‎
Lines changed: 2 additions & 34 deletions b/‎milatools/cli/commands.py‎
Lines changed: 2 additions & 34 deletions
@@ -5,10 +5,10 @@ on:
     branches:
       - master
     paths-ignore:
-      - 'README.md'
+      - "README.md"
   pull_request:
     paths-ignore:
-      - 'README.md'
+      - "README.md"
 
 # https://stackoverflow.com/a/72408109/6388696
 # https://docs.github.com/en/actions/using-jobs/using-concurrency#example-using-concurrency-to-cancel-any-in-progress-job-or-run
@@ -24,7 +24,7 @@ jobs:
       - uses: actions/checkout@v5
       - uses: actions/setup-python@v6
         with:
-          python-version: '3.10'
+          python-version: "3.10"
       - run: pip install "pre-commit<4.0.0"
       - run: pre-commit --version
       - run: pre-commit install
@@ -37,49 +37,50 @@ jobs:
       max-parallel: 4
       matrix:
         platform: [ubuntu-latest, windows-latest, macos-latest]
-        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
     env:
       PLATFORM: ${{ matrix.platform }}
     steps:
-    - uses: actions/checkout@v5
-    - name: Install the latest version of uv
-      uses: astral-sh/setup-uv@v7
-      with:
-        version: "latest"
-        enable-cache: true
-        # https://github.com/astral-sh/setup-uv?tab=readme-ov-file#github-authentication-token
-        github-token: ${{ secrets.GITHUB_TOKEN }}
-        cache-suffix: ${{ matrix.python-version }}
-    - name: Pin Python version to ${{ matrix.python-version }}
-      run: uv python pin ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: uv sync
-    - name: Setup passwordless SSH access to localhost for tests
-      # Adapted from https://stackoverflow.com/a/60367309/6388696
-      if: runner.os == 'Linux'
-      run: |
-        ssh-keygen -t ed25519 -f ~/.ssh/testkey -N ''
-        cat > ~/.ssh/config <<EOF
-          Host localhost
-            User $USER
-            HostName 127.0.0.1
-            IdentityFile ~/.ssh/testkey
-        EOF
-        echo -n 'from="127.0.0.1" ' | cat - ~/.ssh/testkey.pub > ~/.ssh/authorized_keys
-        chmod og-rw ~
-        ssh -o 'StrictHostKeyChecking no' localhost id
+      - uses: actions/checkout@v5
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          version: "latest"
+          enable-cache: true
+          # https://github.com/astral-sh/setup-uv?tab=readme-ov-file#github-authentication-token
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          cache-suffix: ${{ matrix.python-version }}
+      - name: Pin Python version to ${{ matrix.python-version }}
+        run: uv python pin ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: uv sync
+      - name: Setup passwordless SSH access to localhost for tests
+        # Adapted from https://stackoverflow.com/a/60367309/6388696
+        if: runner.os == 'Linux'
+        run: |
+          ssh-keygen -t ed25519 -f ~/_id_ed25519_test -N ''
+          mkdir -p ~/.ssh
+          chmod 700 ~/.ssh
+          cat > ~/.ssh/config <<EOF
+            Host localhost
+              User $USER
+              HostName 127.0.0.1
+              IdentityFile ~/_id_ed25519_test
+          EOF
+          echo -n 'from="127.0.0.1" ' | cat - ~/_id_ed25519_test.pub > ~/.ssh/authorized_keys
+          chmod og-rw ~
+          ssh -o 'StrictHostKeyChecking no' localhost id
 
-    - name: Test with pytest
-      run: uv run pytest --cov=milatools --cov-report=xml --cov-append
+      - name: Test with pytest
+        run: uv run pytest --cov=milatools --cov-report=xml --cov-append
 
-    - name: Store coverage report as an artifact
-      uses: actions/upload-artifact@v4
-      with:
-        name: coverage-reports-unit-${{ matrix.platform }}-${{ matrix.python-version }}
-        path: ./coverage.xml
+      - name: Store coverage report as an artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-reports-unit-${{ matrix.platform }}-${{ matrix.python-version }}
+          path: ./coverage.xml
 
   real-slurm-integration-tests:
-
     name: integration tests with a real SLURM cluster
     needs: [unit-tests]
 
@@ -88,8 +89,8 @@ jobs:
       matrix:
         # TODO: We should ideally also run this with Windows/Mac clients and a Linux
         # server. Unsure how to set that up with GitHub Actions though.
-        python-version: ['3.11']
-        cluster: ['mila']
+        python-version: ["3.11"]
+        cluster: ["mila"]
     uses: ./.github/workflows/testing.yml
     with:
       cluster: ${{ matrix.cluster }}
 
@@ -1,3 +1,10 @@
+"""
+
+TODO: If we need to add a ssh-key on the login node and add it to authorized_keys to get
+ssh access to compute nodes, do it here instead of in mila init, since we're already
+going to go through 2FA once.
+"""
+
 from __future__ import annotations
 
 import asyncio
@@ -7,10 +14,19 @@
 from pathlib import PurePosixPath
 
 from milatools.cli import console
-from milatools.cli.init_command import DRAC_CLUSTERS
+from milatools.cli.init_command import (
+    DEFAULT_DRAC_PUBKEY_PATH,
+    DEFAULT_MILA_PUBKEY_PATH,
+    DRAC_CLUSTERS,
+    can_access_compute_nodes,
+    get_ssh_public_key_path,
+    setup_access_to_compute_nodes,
+)
 from milatools.cli.utils import (
+    SSH_CONFIG_FILE,
     CommandNotFoundError,
     MilatoolsUserError,
+    SSHConfig,
     currently_in_a_test,
     internet_on_compute_nodes,
     running_inside_WSL,
@@ -83,6 +99,15 @@ async def code(
             f"Unable to check the disk-quota on the {cluster} cluster: {exc}"
         )
 
+    public_key_path = get_ssh_public_key_path(
+        cluster, ssh_config=SSHConfig(SSH_CONFIG_FILE)
+    ) or (DEFAULT_MILA_PUBKEY_PATH if cluster == "mila" else DEFAULT_DRAC_PUBKEY_PATH)
+    if not can_access_compute_nodes(login_node, public_key_path=public_key_path):
+        logger.info("Setting up access to compute nodes...")
+        setup_access_to_compute_nodes(
+            cluster, remote=login_node, public_key_path=public_key_path
+        )
+
     # NOTE: Perhaps we could eventually do this check dynamically, if the cluster is an
     # unknown cluster?
     sync_vscode_extensions_task = None
 
@@ -34,12 +34,7 @@
 from milatools.cli import console
 from milatools.cli.code import code
 from milatools.cli.init_command import (
-    print_welcome_message,
-    setup_keys_on_login_node,
-    setup_passwordless_ssh_access,
-    setup_ssh_config,
-    setup_vscode_settings,
-    setup_windows_ssh_config_from_wsl,
+    init,
 )
 from milatools.cli.profile import ensure_program, setup_profile
 from milatools.cli.utils import (
@@ -54,7 +49,6 @@
     get_fully_qualified_name,
     get_hostname_to_use_for_compute_node,
     randname,
-    running_inside_WSL,
     with_control_file,
 )
 from milatools.utils.disk_quota import check_disk_quota_v1
@@ -108,6 +102,7 @@ def main():
         github_issue_url = (
             f"https://github.com/mila-iqia/milatools/issues/new?{urlencode(options)}"
         )
+
         print(
             T.bold_yellow(
                 f"An error occurred during the execution of the command `{command}`. "
@@ -508,33 +503,6 @@ def intranet(search: Sequence[str]) -> None:
     webbrowser.open(url)
 
 
-def init():
-    """Set up your configuration and credentials."""
-
-    #############################
-    # Step 1: SSH Configuration #
-    #############################
-
-    print("Checking ssh config")
-
-    ssh_config = setup_ssh_config()
-
-    success = setup_passwordless_ssh_access(ssh_config=ssh_config)
-    if not success:
-        exit()
-
-    # if we're running on WSL, we actually just copy the id_rsa + id_rsa.pub and the
-    # ~/.ssh/config to the Windows ssh directory (taking care to remove the
-    # ControlMaster-related entries) so that the user doesn't need to install Python on
-    # the Windows side.
-    if running_inside_WSL():
-        setup_windows_ssh_config_from_wsl(linux_ssh_config=ssh_config)
-
-    setup_keys_on_login_node()
-    setup_vscode_settings()
-    print_welcome_message()
-
-
 def forward(
     remote: str,
     page: str | None,