Untitled
unknown
python
3 years ago
8.8 kB
8
Indexable
import os
import json
import paramiko
from scp import SCPClient
from pathlib import *
def to_windows(file_name: str):
return file_name.replace("/", "\\")
def to_linux(file_name: str):
linux_name = file_name.replace("\\", "/")
linux_name = linux_name.replace("/\./", "/")
return linux_name
# =============================================================================
class SSH(object):
hostname = ''
username = ''
password = ''
remote_root = ''
remote_dir = ''
# =============================================================================
def __init__(self, root: str):
defaults_file = os.environ['AT_DEFAULTS']
json_file = os.path.join(os.path.dirname(defaults_file), "cluster.json")
with open(json_file) as f:
auth = json.load(f)
self.hostname = auth["hostname"]
self.username = auth["username"]
self.password = auth["password"]
self.remote_root = root
self.remote_dir = self.remote_root
self.ssh = paramiko.SSHClient()
self.ssh.load_system_host_keys()
self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
# self.ssh.connect(self.hostname, username=self.username, password=self.password, port=22, timeout=3600000000)
self.ssh.connect(self.hostname, username=self.username, password=self.password, port=22)
self.execute("mkdir -p " + self.remote_root)
ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command("mkdir -p " + self.remote_dir)
# =============================================================================
def close(self):
self.ssh.close()
# =============================================================================
def create_remote_dir(self, new_dir: str):
ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command("mkdir -p " + new_dir)
print("Create remote dir " + new_dir)
if ssh_stderr:
print(ssh_stderr)
else:
print("Ok")
# =============================================================================
def change_remote_dir(self, dir_name):
self.remote_dir = self.remote_root + "/" + dir_name # assuming remote is a Linux system
# =============================================================================
def upload(self, file_list, target_dir):
print("************** Upload ***********************")
with SCPClient(self.ssh.get_transport()) as scp:
for f in file_list:
source = to_windows(f.strip())
print("Uploading ", source, "from", os.getcwd(), " to ", target_dir)
scp.put(source, target_dir)
# =============================================================================
def download(self, file_list):
print("************** Download *********************")
# ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command("pwd")
with SCPClient(self.ssh.get_transport()) as scp:
for f in file_list:
source = self.remote_dir + "/" + f
target = os.getcwd()
print("Downloading ", source, " to ", target)
scp.get(source, target)
# =============================================================================
def execute(self, remote_command: str = "", silent: bool = False):
full_command = 'cd ' + self.remote_dir + "; " + remote_command
ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command("pwd")
pwd = ssh_stdout.read().decode().strip()
ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command(full_command)
output = ssh_stdout.read().decode()
error = ssh_stderr.read().decode()
if not silent:
print("************** Execute ***********************")
print(self.username, "@", self.hostname, ":", pwd)
print(full_command)
print("************** Output ***********************")
print(output)
if error:
print("************** Error ************************")
print(error)
# print("************** Done ***********************")
# =============================================================================
class RemoteExecution(object):
LOCAL_BASE_DIR = None
LOCAL_SOURCE_DIR = None
REMOTE_ROOT_DIR = None
REMOTE_BASE_DIR = None
REMOTE_SOURCE_DIR = None
# =========================================================================
def __init__(self, inp_data):
self.LOCAL_BASE_DIR = inp_data["Environment"]["BASE_DIR"]
self.LOCAL_SOURCE_DIR = os.environ["AT_ROOT"]
self.LOCAL_BASE_DIR = self.LOCAL_BASE_DIR.replace("\\", "/")
split = r'proc/'
self.REMOTE_ROOT_DIR = inp_data["RemoteServer"]["root_dir"]
try:
self.REMOTE_BASE_DIR = os.path.join(self.REMOTE_ROOT_DIR, self.LOCAL_BASE_DIR.split(split)[1])
except IndexError:
self.REMOTE_BASE_DIR = os.path.join(self.REMOTE_ROOT_DIR, os.path.basename(self.LOCAL_BASE_DIR))
self.REMOTE_BASE_DIR = self.REMOTE_BASE_DIR.replace("\\", "/")
print("Remote Base Dir: ", self.REMOTE_BASE_DIR)
self.REMOTE_SOURCE_DIR = inp_data["RemoteServer"]["source_dir"]
self.source_file_filters = inp_data["RemoteServer"]["source_file_filter"]
# =============================================================================
def init_remote(self):
print("Initializing remote session")
uploads = [str(x) for x in Path(".").iterdir() if x.is_file()]
print("DBG hhh: Remote Base Dir: ", self.REMOTE_BASE_DIR)
ssh = SSH(self.REMOTE_BASE_DIR)
ssh.upload(uploads, self.REMOTE_BASE_DIR)
ssh.close()
# =============================================================================
def upload_sources(self):
print("Uploading sources")
os.chdir(self.LOCAL_SOURCE_DIR)
ssh = SSH(self.REMOTE_SOURCE_DIR)
for file_filter in self.source_file_filters:
for root, dirs, files in os.walk("."):
for file in files:
if file.endswith(".py"):
if "*" in file_filter or file_filter in file:
source = os.path.join(root, file)
target = os.path.join(self.REMOTE_SOURCE_DIR, source)
target_dir = str(Path(target).parent).replace("\\", "/")
ssh.upload([source], target_dir)
ssh.close()
os.chdir(self.LOCAL_BASE_DIR)
# =============================================================================
def run_remote(self, step_name, step_dir, uploads, downloads, upload_all=False):
print("Running step_name " + step_name + " remotely")
REMOTE_STEP_DIR = os.path.join(self.REMOTE_BASE_DIR, step_dir).replace("\\", "/")
ssh = SSH(self.REMOTE_BASE_DIR)
ssh.create_remote_dir(REMOTE_STEP_DIR)
# ssh.upload(["../"+self.REPORT_FILE], self.REMOTE_BASE_DIR)
if upload_all:
uploads = [str(x) for x in Path(".").iterdir() if x.is_file()]
ssh.upload(uploads, REMOTE_STEP_DIR)
if step_name == "ALL":
ssh.execute("python $AT_ROOT/runCase.py --nogui >& runCase.log")
else:
ssh.execute("python $AT_ROOT/runCase.py --nogui --slave --input_file runCase.inp.yml --step=" + step_name + " >& " + step_name + ".log")
ssh.change_remote_dir(step_dir)
ssh.download(downloads)
ssh.close()
# =============================================================================
def download_file(self, file):
ssh = SSH(self.REMOTE_BASE_DIR)
ssh.download([file])
# ====================================================
# Template for Main
# ====================================================
if __name__ == "__main__":
from Lib.user_arguments import UserOptions
usr = UserOptions()
usr.input_file.EXTENSION = "*.inp"
usr.input_file.DESCRIPTION = "Input file [runCase.inp] "
usr.input_file.REQUIRED = True
usr.input_file.DEFAULT = "runCase.inp"
usr.work_file.REQUIRED = False
usr.work_file.ENABLE_GUI = False
usr.get_options()
inp_file = usr.input_file.NAME
os.chdir(usr.input_file.DIRECTORY)
rem = SSH()
rem.change_remote_dir("Transfer_speed")
rem.upload([usr.input_file.BASENAME])
rem.execute("cp geom1.inp mima.inp")
rem.execute("ls")
rem.download(["mima.inp"])
Editor is loading...