def run(self):
# copy input to remote host
# Remote targets need to be copied first to local system
inp = self.input()
if isinstance(inp, RemoteTarget):
inp.get(inp.path)
remote_fs = RemoteFileSystem(self.ssh_host)
remote_path = os.path.join(self.remote_dir, self.input().path)
remote_fs.put(self.input().path, remote_path)
infile = self.input().path
outfile = os.path.basename(self.output().path)
remote = RemoteContext(self.ssh_host)
remote.check_output(
['cd {};'.format(self.remote_dir), self.exe, infile, outfile])
# clean up intermediate files
if isinstance(inp, RemoteTarget):
os.remove(infile)
remote_fs.remove(remote_path)
if self.remove_input:
self.input().remove()
def setUp(self):
self.ctx = RemoteContext(working_ssh_host)
self.filepath = "/tmp/luigi_remote_test.dat"
self.target = RemoteTarget(
self.filepath,
working_ssh_host,
)
self.ctx.check_output(["rm", "-rf", self.filepath])
self.ctx.check_output(["echo -n 'hello' >", self.filepath])
def run(self):
# copy input to remote host
# Remote targets need to be copied first to local system
inp = self.input()
if isinstance(inp, RemoteTarget):
inp.get(inp.path)
remote_fs = RemoteFileSystem(self.ssh_host)
remote_path = os.path.join(self.remote_dir, self.input().path)
remote_fs.put(self.input().path, remote_path)
infile = self.input().path
outfile = os.path.basename(self.output().path)
# create slurm batch script
with open('slurm_job.sh', 'w') as f:
f.write('#!/bin/bash\n')
f.write('#SBATCH --job-name=luigi-task\n')
f.write('#SBATCH --time={}\n'.format(self.slurm_time))
f.write('#SBATCH --partition={}\n'.format(self.slurm_partition))
f.write('#SBATCH --ntasks=1\n')
f.write('#SBATCH --account=project_2002078\n')
f.write('\n')
f.write('srun {} {} {}\n'.format(self.exe, infile, outfile))
f.write('\n')
# copy batch script to remote host
remote_path = os.path.join(self.remote_dir, 'slurm_job.sh')
remote_fs.put('slurm_job.sh', remote_path)
os.remove('slurm_job.sh')
# Submit batch job and wait for its completion
remote = RemoteContext(self.ssh_host)
slurm_out = remote.check_output(
['cd {};'.format(self.remote_dir), 'sbatch', 'slurm_job.sh'])
job_id = int(slurm_out.split()[-1])
sacct_cmd = 'sacct -n -X -j {}'.format(job_id)
# Wait for job to complete
sleep(20)
while True:
sacct_out = remote.check_output([sacct_cmd]).decode('utf-8')
if 'COMPLETED' in sacct_out:
print("DEB: completed")
break
elif 'FAILED' in sacct_out:
self.output().remove()
raise RuntimeError('Task {} failed'.format(
self.__class__.__name__))
else:
# Do not poll sacct too frequently
sleep(60)
def run(self):
remote = RemoteContext(self.document_server)
with self.output().open('w') as outfile:
stream = StringIO(
remote.check_output([
"cd {} && find ./{} -type f -regex \".*[w]*arc\.gz\"".
format(self.data_collections_folder, self.collection_name)
]).decode())
for line in stream:
outfile.write(
line.replace(
'./', 'http://{}:8080/browser/files/'.format(
self.document_server)))
def test_check_output_fail_connect(self):
""" Test check_output to a non-existing host """
context = RemoteContext("__NO_HOST_LIKE_THIS__", connect_timeout=1)
self.assertRaises(
subprocess.CalledProcessError,
context.check_output, ["ls"]
)
def return_ssh_context(host: str, username: str, key_file: str):
try:
return RemoteContext(host=host,
username=username,
key_file=key_file,
no_host_key_check=True)
except RemoteCalledProcessError as e:
print(f"context error: {e.__str__()}")
class TestRemoteContext(unittest.TestCase):
def setUp(self):
self.context = RemoteContext(working_ssh_host)
def tearDown(self):
try:
self.remote_server_handle.terminate()
except Exception:
pass
def test_check_output(self):
""" Test check_output ssh
Assumes the running user can ssh to working_ssh_host
"""
output = self.context.check_output(["echo", "-n", "luigi"])
self.assertEqual(output, b"luigi")
def test_tunnel(self):
print("Setting up remote listener...")
self.remote_server_handle = self.context.Popen([
"python", "-c", '"{0}"'.format(HELLO_SERVER_CMD)
], stdout=subprocess.PIPE)
print("Setting up tunnel")
with self.context.tunnel(2135, 2134):
print("Tunnel up!")
# hack to make sure the listener process is up
# and running before we write to it
server_output = self.remote_server_handle.stdout.read(5)
self.assertEqual(server_output, b"ready")
print("Connecting to server via tunnel")
s = socket.socket()
s.connect(("localhost", 2135))
print("Receiving...",)
response = s.recv(5)
self.assertEqual(response, b"hello")
print("Closing connection")
s.close()
print("Waiting for listener...")
output, _ = self.remote_server_handle.communicate()
self.assertEqual(self.remote_server_handle.returncode, 0)
print("Closing tunnel")
class TestRemoteContext(unittest.TestCase):
def setUp(self):
self.context = RemoteContext(working_ssh_host)
def tearDown(self):
try:
self.remote_server_handle.terminate()
except Exception:
pass
def test_check_output(self):
""" Test check_output ssh
Assumes the running user can ssh to working_ssh_host
"""
output = self.context.check_output(["echo", "-n", "luigi"])
self.assertEqual(output, b"luigi")
def test_tunnel(self):
print("Setting up remote listener...")
self.remote_server_handle = self.context.Popen([
"python", "-c", '"{0}"'.format(HELLO_SERVER_CMD)
], stdout=subprocess.PIPE)
print("Setting up tunnel")
with self.context.tunnel(2135, 2134):
print("Tunnel up!")
# hack to make sure the listener process is up
# and running before we write to it
server_output = self.remote_server_handle.stdout.read(5)
self.assertEqual(server_output, b"ready")
print("Connecting to server via tunnel")
s = socket.socket()
s.connect(("localhost", 2135))
print("Receiving...",)
response = s.recv(5)
self.assertEqual(response, b"hello")
print("Closing connection")
s.close()
print("Waiting for listener...")
output, _ = self.remote_server_handle.communicate()
self.assertEqual(self.remote_server_handle.returncode, 0)
print("Closing tunnel")
def setUp(self):
self.ctx = RemoteContext(working_ssh_host)
self.filepath = "/tmp/luigi_remote_test.dat"
self.target = RemoteTarget(
self.filepath,
working_ssh_host,
)
self.ctx.check_output(["rm", "-rf", self.filepath])
self.ctx.check_output(["echo -n 'hello' >", self.filepath])
def test_subprocess_delegation(self):
""" Test subprocess call structure using mock module """
orig_Popen = subprocess.Popen
self.last_test = None
def Popen(cmd, **kwargs):
self.last_test = cmd
subprocess.Popen = Popen
context = RemoteContext("some_host",
username="******",
key_file="/some/key.pub")
context.Popen(["ls"])
self.assertTrue("ssh" in self.last_test)
self.assertTrue("-i" in self.last_test)
self.assertTrue("/some/key.pub" in self.last_test)
self.assertTrue("luigi@some_host" in self.last_test)
self.assertTrue("ls" in self.last_test)
subprocess.Popen = orig_Popen
class TestRemoteTargetAtomicity(unittest.TestCase):
path = '/tmp/luigi_remote_atomic_test.txt'
ctx = RemoteContext(working_ssh_host)
def _exists(self, path):
try:
self.ctx.check_output(["test", "-e", path])
except subprocess.CalledProcessError, e:
if e.returncode == 1:
return False
else:
raise
return True
class TestRemoteTargetAtomicity(unittest.TestCase,
target_test.FileSystemTargetTestMixin):
path = '/tmp/luigi_remote_atomic_test.txt'
ctx = RemoteContext(working_ssh_host)
def create_target(self, format=None):
return RemoteTarget(self.path, working_ssh_host, format=format)
def _exists(self, path):
try:
self.ctx.check_output(["test", "-e", path])
except subprocess.CalledProcessError as e:
if e.returncode == 1:
return False
else:
raise
return True
def assertCleanUp(self, tp):
self.assertFalse(self._exists(tp))
def setUp(self):
self.ctx.check_output(["rm", "-rf", self.path])
self.local_file = '/tmp/local_luigi_remote_atomic_test.txt'
if os.path.exists(self.local_file):
os.remove(self.local_file)
def tearDown(self):
self.ctx.check_output(["rm", "-rf", self.path])
if os.path.exists(self.local_file):
os.remove(self.local_file)
def test_put(self):
f = open(self.local_file, 'w')
f.write('hello')
f.close()
t = RemoteTarget(self.path, working_ssh_host)
t.put(self.local_file)
self.assertTrue(self._exists(self.path))
def test_get(self):
self.ctx.check_output(["echo -n 'hello' >", self.path])
t = RemoteTarget(self.path, working_ssh_host)
t.get(self.local_file)
f = open(self.local_file, 'r')
file_content = f.read()
self.assertEqual(file_content, 'hello')
test_move_on_fs = None # ssh don't have move (yet?)
test_rename_dont_move_on_fs = None # ssh don't have move (yet?)
class TestRemoteTarget(unittest.TestCase):
""" These tests assume RemoteContext working
in order for setUp and tearDown to work
"""
def setUp(self):
self.ctx = RemoteContext(working_ssh_host)
self.filepath = "/tmp/luigi_remote_test.dat"
self.target = RemoteTarget(
self.filepath,
working_ssh_host,
)
self.ctx.check_output(["rm", "-rf", self.filepath])
self.ctx.check_output(["echo -n 'hello' >", self.filepath])
def tearDown(self):
self.ctx.check_output(["rm", "-rf", self.filepath])
def test_exists(self):
self.assertTrue(self.target.exists())
no_file = RemoteTarget(
"/tmp/_file_that_doesnt_exist_",
working_ssh_host,
)
self.assertFalse(no_file.exists())
def test_remove(self):
self.target.remove()
self.assertRaises(
subprocess.CalledProcessError,
self.ctx.check_output,
["cat", self.filepath]
)
def test_open(self):
f = self.target.open('r')
file_content = f.read()
f.close()
self.assertEqual(file_content, "hello")
self.assertTrue(self.target.fs.exists(self.filepath))
self.assertFalse(self.target.fs.isdir(self.filepath))
def test_context_manager(self):
with self.target.open('r') as f:
file_content = f.read()
self.assertEqual(file_content, "hello")
class TestRemoteTarget(unittest.TestCase):
""" These tests assume RemoteContext working
in order for setUp and tearDown to work
"""
def setUp(self):
self.ctx = RemoteContext(working_ssh_host)
self.filepath = "/tmp/luigi_remote_test.dat"
self.target = RemoteTarget(
self.filepath,
working_ssh_host,
)
self.ctx.check_output(["rm", "-rf", self.filepath])
self.ctx.check_output(["echo -n 'hello' >", self.filepath])
def tearDown(self):
self.ctx.check_output(["rm", "-rf", self.filepath])
def test_exists(self):
self.assertTrue(self.target.exists())
no_file = RemoteTarget(
"/tmp/_file_that_doesnt_exist_",
working_ssh_host,
)
self.assertFalse(no_file.exists())
def test_remove(self):
self.target.remove()
self.assertRaises(
subprocess.CalledProcessError,
self.ctx.check_output,
["cat", self.filepath]
)
def test_open(self):
f = self.target.open('r')
file_content = f.read()
f.close()
self.assertEqual(file_content, "hello")
self.assertTrue(self.target.fs.exists(self.filepath))
self.assertFalse(self.target.fs.isdir(self.filepath))
def test_context_manager(self):
with self.target.open('r') as f:
file_content = f.read()
self.assertEqual(file_content, "hello")
def execute_command(context: RemoteContext, command: str):
try:
output = context.check_output(shlex.split(command))
return output
except RemoteCalledProcessError as e:
print(f"execute error: {e.__str__()}")
def setUp(self):
self.context = RemoteContext(working_ssh_host)
def run(self):
remote = RemoteContext(SSH_HOST)
print remote.check_output([
"ps aux > {0}".format(self.output().path)
])
def run(self):
parameters.setParameters(parameters.data)
remote = RemoteContext(SSH_HOST)
remote.check_output(['mkdir -p ' + parameters.data['remote_storage']])
def run(self):
remote = RemoteContext(SSH_HOST)
print remote.check_output(["ps aux > {0}".format(self.output().path)])
class TestRemoteTargetAtomicity(unittest.TestCase):
path = '/tmp/luigi_remote_atomic_test.txt'
ctx = RemoteContext(working_ssh_host)
def _exists(self, path):
try:
self.ctx.check_output(["test", "-e", path])
except subprocess.CalledProcessError as e:
if e.returncode == 1:
return False
else:
raise
return True
def setUp(self):
self.ctx.check_output(["rm", "-rf", self.path])
self.local_file = '/tmp/local_luigi_remote_atomic_test.txt'
if os.path.exists(self.local_file):
os.remove(self.local_file)
def tearDown(self):
self.ctx.check_output(["rm", "-rf", self.path])
if os.path.exists(self.local_file):
os.remove(self.local_file)
def test_close(self):
t = RemoteTarget(self.path, working_ssh_host)
p = t.open('w')
print('test', file=p)
self.assertFalse(self._exists(self.path))
p.close()
self.assertTrue(self._exists(self.path))
def test_del(self):
t = RemoteTarget(self.path, working_ssh_host)
p = t.open('w')
print('test', file=p)
tp = p.tmp_path
del p
self.assertFalse(self._exists(tp))
self.assertFalse(self._exists(self.path))
def test_write_cleanup_no_close(self):
t = RemoteTarget(self.path, working_ssh_host)
def context():
f = t.open('w')
f.write('stuff')
context()
gc.collect() # force garbage collection of f variable
self.assertFalse(t.exists())
def test_write_cleanup_with_error(self):
t = RemoteTarget(self.path, working_ssh_host)
try:
with t.open('w'):
raise Exception('something broke')
except:
pass
self.assertFalse(t.exists())
def test_write_with_success(self):
t = RemoteTarget(self.path, working_ssh_host)
with t.open('w') as p:
p.write("hello")
self.assertTrue(t.exists())
def test_gzip(self):
t = RemoteTarget(self.path, working_ssh_host, luigi.format.Gzip)
p = t.open('w')
test_data = 'test'
p.write(test_data)
self.assertFalse(self._exists(self.path))
p.close()
self.assertTrue(self._exists(self.path))
# Using gzip module as validation
cmd = 'scp -q %s:%s %s' % (working_ssh_host, self.path,
self.local_file)
assert os.system(cmd) == 0
f = gzip.open(self.local_file, 'rb')
self.assertTrue(test_data == f.read())
f.close()
# Verifying our own gzip remote reader
f = RemoteTarget(self.path, working_ssh_host,
luigi.format.Gzip).open('r')
self.assertTrue(test_data == f.read())
f.close()
def test_put(self):
f = open(self.local_file, 'w')
f.write('hello')
f.close()
t = RemoteTarget(self.path, working_ssh_host)
t.put(self.local_file)
self.assertTrue(self._exists(self.path))
def test_get(self):
self.ctx.check_output(["echo -n 'hello' >", self.path])
t = RemoteTarget(self.path, working_ssh_host)
t.get(self.local_file)
f = open(self.local_file, 'r')
file_content = f.read()
self.assertEqual(file_content, 'hello')
def setUp(self):
self.context = RemoteContext(working_ssh_host)
