7d11dfa688
Implementation of a mock scheduler for jobs/tasks that will evolve into a job manager for running training jobs over vast.ai instances.
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Protocol, runtime_checkable
|
|
|
|
|
|
@runtime_checkable
|
|
class Connection(Protocol):
|
|
async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]: ...
|
|
async def push(self, local_path: str, remote_path: str) -> None: ...
|
|
async def pull(self, remote_path: str, local_path: str) -> None: ...
|
|
async def close(self) -> None: ...
|
|
|
|
|
|
@runtime_checkable
|
|
class Transport(Protocol):
|
|
async def connect(self, host: str, port: int) -> Connection: ...
|
|
|
|
|
|
class SSHConnection:
|
|
def __init__(self, host: str, port: int):
|
|
self.host = host
|
|
self.port = port
|
|
|
|
async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"ssh",
|
|
"-p", str(self.port),
|
|
"-o", "StrictHostKeyChecking=no",
|
|
"-o", "ConnectTimeout=5",
|
|
"-o", "BatchMode=yes",
|
|
f"root@{self.host}",
|
|
cmd,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
try:
|
|
stdout, stderr = await asyncio.wait_for(
|
|
proc.communicate(), timeout=timeout
|
|
)
|
|
return proc.returncode, (stdout + stderr).decode()
|
|
except asyncio.TimeoutError:
|
|
proc.kill()
|
|
return 1, "timeout"
|
|
|
|
async def push(self, local_path: str, remote_path: str) -> None:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"scp",
|
|
"-P", str(self.port),
|
|
"-o", "StrictHostKeyChecking=no",
|
|
local_path,
|
|
f"root@{self.host}:{remote_path}",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
_, stderr = await proc.communicate()
|
|
if proc.returncode != 0:
|
|
raise ConnectionError(
|
|
f"scp push failed to {self.host}:{self.port}: {stderr.decode()}"
|
|
)
|
|
|
|
async def pull(self, remote_path: str, local_path: str) -> None:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"scp",
|
|
"-P", str(self.port),
|
|
"-o", "StrictHostKeyChecking=no",
|
|
f"root@{self.host}:{remote_path}",
|
|
local_path,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
_, stderr = await proc.communicate()
|
|
if proc.returncode != 0:
|
|
raise ConnectionError(
|
|
f"scp pull failed from {self.host}:{self.port}: {stderr.decode()}"
|
|
)
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
|
|
class SSHTransport:
|
|
async def connect(self, host: str, port: int) -> SSHConnection:
|
|
return SSHConnection(host, port)
|