Transition `mila code` to RemoteV2 [MT-80] #105

lebrice · 2024-03-21T21:30:43Z

Based off of #112

Support ssh keys with passphrase in mila code (thanks to RemoteV2)
Adjust code to use RemoteV2.
Rename code to code_v1, code --> code_v1 for Windows and add a new code implementation for WSL/MacOS/Linux.
- For now, leaving the current mila code command for Windows users, printing a warning encouraging to switch to WSL.

codecov-commenter · 2024-03-26T17:52:10Z

Codecov Report

Attention: Patch coverage is 85.67335% with 50 lines in your changes are missing coverage. Please review.

Project coverage is 74.55%. Comparing base (9b4a7cd) to head (1972cf8).

Files	Patch %	Lines
milatools/utils/vscode_utils.py	80.21%	18 Missing ⚠️
milatools/cli/code.py	82.43%	13 Missing ⚠️
milatools/utils/disk_quota.py	86.95%	9 Missing ⚠️
milatools/utils/parallel_progress.py	88.13%	7 Missing ⚠️
milatools/cli/commands.py	92.50%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #105      +/-   ##
==========================================
+ Coverage   72.93%   74.55%   +1.61%     
==========================================
  Files          15       17       +2     
  Lines        2176     2283     +107     
==========================================
+ Hits         1587     1702     +115     
+ Misses        589      581       -8

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Signed-off-by: Fabrice Normandin <[email protected]>

satyaog

Code review done, will do the tests review

satyaog · 2024-05-29T16:59:41Z

milatools/cli/code.py

+    """
+    # Check that the `code` command is in the $PATH so that we can use just `code` as
+    # the command.
+    code_command = command


do we need code_command? Can we just use command?

Sure. Fixed in 4feafe0

satyaog · 2024-05-29T21:25:14Z

milatools/cli/code.py

+        # todo: use the mila or the local machine as the reference for vscode
+        # extensions?


I think I would say mila but I'm not sure if this would change much as users would probably use the sync feature of vscode it self? I wonder if there could be incompatibility between local and cluster if local is not linux?

Not sure either! I guess we'll have to test it out.

milatools/cli/code.py

satyaog · 2024-05-29T22:23:07Z

milatools/cli/code.py

+        # ideally, so we can collect some simple stats about the use of `milatools` on
+        # the clusters.
+        if any(flag == "-J" or "-J=" in flag or "--job-name" in flag for flag in alloc):
+            # todo: Get the job name from the flags instead?


A warning could be a good idea with something like this :

Suggested change

# todo: Get the job name from the flags instead?

import argparse

parser = argparse.ArgumentParser()

parser.add_argument("-J", "--job-name")

args, alloc = parser.parse_known_args(alloc)

if args.job_name:

#warning

I'm not parsing the job name at the moment, but yes I'd definitely do something like this if we decide to do so.

satyaog · 2024-05-29T22:46:46Z

milatools/cli/code.py

+        # TODO: BUG: This now requires two Ctrl+C's instead of one!
+        console.print(
+            "The editor was closed. Reopen it with <Enter> or terminate the "
+            "process with <Ctrl+C> (maybe twice)."
+        )


Do we know why?

This had to do with how asyncio deals with KeyboardInterrupt, but I might have fixed it. I'll double check.

Yeah it still seems necessary to do Ctrl+C twice. I'm not sure why exactly.

satyaog · 2024-05-30T00:29:29Z

milatools/utils/parallel_progress.py

+    if not currently_in_a_test():
+        columns.append(TimeElapsedColumn())


This should probably be mocked to output a fixed string instead of mixing test code to real code

Fixed in 31b55dd

satyaog · 2024-05-30T00:31:07Z

milatools/utils/parallel_progress.py

+    )
+    # iterate over the jobs we need to run
+    for task_description, async_task_fn in zip(task_descriptions, async_task_fns):
+        # NOTE: Could set visible=false so we don't have a lot of bars all at once.


+1 for the idea

Applied in 94f690d

satyaog · 2024-05-30T00:35:34Z

milatools/utils/parallel_progress.py


-    for result in parallel_progress_bar([_example_task_fn, _example_task_fn]):
-        print(result)
+class AsyncTaskFn(Protocol[OutT_co]):


Is this for typing help cues or does this have a practical utility in the async progress bars?

This is the signature of a function that is used by the async progress bar.

your IDE is then able to show you if your function is valid by checking against this protocol

satyaog · 2024-05-30T14:13:46Z

milatools/utils/vscode_utils.py

+    # Connect to the remotes in parallel.
+    dest_runners_and_hostnames = await asyncio.gather(
+        *(_get_runner_and_hostname(dest) for dest in destinations)
+    )


It does not seam necessary to have this bit of coda and the _get_runner_and_hostname util as _install_vscode_extensions_task_function will connect in parallel anyway?

Fixed in 94f690d

satyaog · 2024-05-30T14:21:20Z

milatools/utils/vscode_utils.py

        if dest_hostname == "localhost":
-            remote = LocalV1()
+            remote = LocalV2()


I'm not sure to understand how and why we would be syncing to localhost? Wouldn't it be simpler to manage the extensions locally through the vscode interface?

You might want to get the same vscode extensions locally as you had on the Mila cluster. Sure, if you have "settings sync" setup in vscode, it could probably sync it better than this would, but I don't see why we shouldn't support doing this.

satyaog · 2024-05-30T14:43:00Z

tests/cli/test_utils.py

+    ],
+)
+def test_internet_on_compute_nodes(cluster: str, expected: bool):
+    assert internet_on_compute_nodes(cluster) == expected


We might want to make a real internet connection test when we'll move away from hardcoded list of clusters

satyaog · 2024-05-30T19:48:44Z