Skip to content

Commit

Permalink
Restore code_v1 to commands.py to reduce changes
Browse files Browse the repository at this point in the history 
Signed-off-by: Fabrice Normandin <[email protected]>
  • Loading branch information
@lebrice
lebrice committed May 16, 2024
1 parent 522f076 commit ef2634f
Show file tree
Hide file tree
Showing 4 changed files with 184 additions and 184 deletions.
177 changes: 5 additions & 172 deletions milatools/cli/code.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,29 +7,23 @@
from logging import getLogger as get_logger
from typing import Awaitable

from typing_extensions import deprecated

from milatools.cli import console
from milatools.cli.commands import find_allocation
from milatools.cli.init_command import DRAC_CLUSTERS
from milatools.cli.utils import (
CLUSTERS,
Cluster,
CommandNotFoundError,
MilatoolsUserError,
currently_in_a_test,
get_hostname_to_use_for_compute_node,
no_internet_on_compute_nodes,
running_inside_WSL,
internet_on_compute_nodes,
)
from milatools.utils.compute_node import (
ComputeNode,
salloc,
sbatch,
)
from milatools.utils.disk_quota import check_disk_quota, check_disk_quota_v1
from milatools.utils.disk_quota import check_disk_quota
from milatools.utils.local_v2 import LocalV2
from milatools.utils.remote_v1 import RemoteV1
from milatools.utils.remote_v2 import (
RemoteV2,
)
Expand Down Expand Up @@ -99,6 +93,8 @@ def add_mila_code_arguments(subparsers: argparse._SubParsersAction):
help="When set, skips the syncing of vscode extensions.",
)
if sys.platform == "win32":
from milatools.cli.commands import code_v1

code_parser.set_defaults(function=code_v1)
else:
code_parser.set_defaults(function=code)
Expand Down Expand Up @@ -152,7 +148,7 @@ async def code(
# NOTE: Perhaps we could eventually do this check dynamically, if the cluster is an
# unknown cluster?
sync_vscode_extensions_task = None
if no_internet_on_compute_nodes(cluster) and not no_sync:
if not internet_on_compute_nodes(cluster) and not no_sync:
# Sync the VsCode extensions from the local machine over to the target cluster.
console.log(
f"Installing VSCode extensions that are on the local machine on "
Expand Down Expand Up @@ -276,166 +272,3 @@ async def launch_vscode_loop(code_command: str, compute_node: ComputeNode, path:
except asyncio.CancelledError:
raise
# return


@deprecated(
"Support for the `mila code` command is now deprecated on Windows machines, as it "
"does not support ssh keys with passphrases or clusters where 2FA is enabled. "
"Please consider switching to the Windows Subsystem for Linux (WSL) to run "
"`mila code`."
)
def code_v1(
path: str,
command: str,
persist: bool,
job: int | None,
node: str | None,
alloc: list[str],
cluster: Cluster = "mila",
no_sync: bool = False,
):
"""Open a remote VSCode session on a compute node.
Arguments:
path: Path to open on the remote machine
command: Command to use to start vscode
(defaults to "code" or the value of $MILATOOLS_CODE_COMMAND)
persist: Whether the server should persist or not
job: Job ID to connect to
node: Node to connect to
alloc: Extra options to pass to slurm
"""
here = LocalV2()
remote = RemoteV1(cluster)

if cluster != "mila" and job is None and node is None:
if not any("--account" in flag for flag in alloc):
logger.warning(
"Warning: When using the DRAC clusters, you usually need to "
"specify the account to use when submitting a job. You can specify "
"this in the job resources with `--alloc`, like so: "
"`--alloc --account=<account_to_use>`, for example:\n"
f"mila code {path} --cluster {cluster} --alloc "
f"--account=your-account-here"
)

if command is None:
command = get_code_command()

try:
check_disk_quota_v1(remote)
except MilatoolsUserError:
raise
except Exception as exc:
logger.warning(f"Unable to check the disk-quota on the cluster: {exc}")

if sys.platform == "win32":
print(
"Syncing vscode extensions in the background isn't supported on "
"Windows. Skipping."
)
elif no_internet_on_compute_nodes(cluster):
# Sync the VsCode extensions from the local machine over to the target cluster.
# TODO: Make this happen in the background (without overwriting the output).
run_in_the_background = False
print(
console.log(
f"[cyan]Installing VSCode extensions that are on the local machine on "
f"{cluster}" + (" in the background." if run_in_the_background else ".")
)
)
asyncio.run(
sync_vscode_extensions(
LocalV2(),
[cluster],
)
)

if node is None:
cnode = find_allocation(
remote,
job_name="mila-code",
job=job,
node=node,
alloc=alloc,
cluster=cluster,
)
if persist:
cnode = cnode.persist()

data, proc = cnode.ensure_allocation()

node_name = data["node_name"]
else:
node_name = node
proc = None
data = None

if not path.startswith("/"):
# Get $HOME because we have to give the full path to code
home = remote.home()
path = home if path == "." else f"{home}/{path}"

command_path = shutil.which(command)
if not command_path:
raise CommandNotFoundError(command)

# NOTE: Since we have the config entries for the DRAC compute nodes, there is no
# need to use the fully qualified hostname here.
if cluster == "mila":
node_name = get_hostname_to_use_for_compute_node(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.
inside_WSL = running_inside_WSL()
try:
while True:
if inside_WSL:
here.run(
(
"powershell.exe",
"code",
"-nw",
"--remote",
f"ssh-remote+{node_name}",
path,
),
)
else:
here.run(
(
command_path,
"-nw",
"--remote",
f"ssh-remote+{node_name}",
path,
),
)
print(
"The editor was closed. Reopen it with <Enter>"
" or terminate the process with <Ctrl+C>"
)
if currently_in_a_test():
break
input()

except KeyboardInterrupt:
if not persist:
if proc is not None:
proc.kill()
print(f"Ended session on '{node_name}'")

if persist:
print("This allocation is persistent and is still active.")
print("To reconnect to this node:")
console.print(
f" mila code {path} "
+ (f"--cluster={cluster} " if cluster != "mila" else "")
+ f"--node {node_name}",
style="bold",
)
print("To kill this allocation:")
assert data is not None
assert "jobid" in data
console.print(f" ssh {cluster} scancel {data['jobid']}", style="bold")
Loading

0 comments on commit ef2634f

Please sign in to comment.