1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
|
from pathlib import Path
from typing import Protocol, Self
from execution import run
from log import info
class Context(Protocol):
def context(self) -> str: ...
async def sync_to(self, experiment_root: Path) -> None: ...
async def run(self, cmd: str | list[str]) -> None: ...
async def sync_back(self, remote_path: Path, local_path: Path) -> None: ...
async def __aenter__(self) -> Self: ...
async def __aexit__(self, exc_type, exc_value, exc_tb): ...
class RemoteContext:
__remote_connection: str
__jump_connection: str | None
__cwd: str | None
__delete: bool
def __init__(
self,
remote_user: str,
remote_host: str,
jump_user: str | None = None,
jump_host: str | None = None,
jump_port: str | None = None,
delete: bool = False,
) -> None:
self.__cwd = None
self.__delete = delete
self.__remote_connection = f"{remote_user}@{remote_host}"
self.__jump_connection = None
if any([jump_host, jump_user, jump_port]):
assert all([jump_host, jump_user]), "cannot define jump partially"
self.__jump_connection = f"ssh -J {jump_user}@{jump_host}"
if jump_port:
self.__jump_connection += f" -p {jump_port}"
def context(self) -> str:
return f"Remote Context: {self.__cwd}"
async def sync_to(self, experiment_root: Path) -> None:
cmd = "rsync -avz"
if self.__jump_connection:
cmd += f" -e '{self.__jump_connection}'"
cmd += f" {experiment_root} {self.__remote_connection}:{self.__cwd}"
info(f"cmd: {cmd}")
_ = await run(cmd)
def __connection_cmd(self) -> str:
connection_cmd = f"ssh {self.__remote_connection}"
if self.__jump_connection:
connection_cmd = f"{self.__jump_connection} {self.__remote_connection}"
return connection_cmd
async def run(self, cmd: str | list[str]) -> None:
assert self.__cwd, "Cannot run command without setup"
if isinstance(cmd, list):
cmd = " && ".join(cmd)
cmd = cmd.strip()
assert cmd.startswith("&&") == False, "already prefixed with '&&'"
_ = await run([f"{self.__connection_cmd()} 'cd {self.__cwd} && {cmd}'"])
async def sync_back(self, remote_path: Path, local_path: Path) -> None:
cmd = "rsync -avz"
if self.__jump_connection:
cmd += f" -e '{self.__jump_connection}'"
cmd += f"{self.__remote_connection}:{self.__cwd}/{remote_path} {local_path}"
_ = await run(cmd)
async def __aenter__(self) -> Self:
info("__aenter__")
directory = await run(f"{self.__connection_cmd()} 'mktemp -d'")
assert directory, "failed to create directory on remote system"
self.__cwd = directory
return self
async def __aexit__(self, exc_type, exc_value, exc_tb):
assert self.__cwd, "incorrect state, cannot exit context which was not entered"
if self.__delete:
_ = await run(f"{self.__connection_cmd()} 'rm -rf {self.__cwd}'")
class LocalContext:
__cwd: str | None
__delete: bool
def __init__(
self,
delete: bool = False,
) -> None:
self.__cwd = None
self.__delete = delete
def context(self) -> str:
return f"Local Context: {self.__cwd}"
async def sync_to(self, experiment_root: Path) -> None:
_ = await run(f"rsync -av {experiment_root} {self.cwd}")
async def run(self, cmd: str | list[str]) -> None:
assert self.cwd, "Cannot run command without setup"
if isinstance(cmd, list):
cmd = " && ".join(cmd)
cmd = cmd.strip()
assert cmd.startswith("&&") == False, "already prefixed with '&&'"
_ = await run([f"cd {self.cwd} && {cmd}"])
async def sync_back(self, remote_path: Path, local_path: Path) -> None:
_ = await run(f"rsync -av {self.cwd}/{remote_path} {local_path}")
async def __aenter__(self) -> Self:
directory = await run("mktemp -d")
assert directory, "failed to create directory on remote system"
self.cwd = directory
return self
async def __aexit__(self, exc_type, exc_value, exc_tb):
assert self.cwd, "incorrect state, cannot exit context which was not entered"
if self.__delete:
_ = await run(f"rm -rf {self.__cwd}")
|