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()
Example #2
0
 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)
Example #4
0
    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)))
Example #5
0
 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__()}")
Example #7
0
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")
Example #8
0
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")
Example #9
0
 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])
Example #10
0
    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
Example #11
0
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
Example #12
0
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?)
Example #13
0
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")
Example #14
0
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__()}")
Example #16
0
 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)
     ])
Example #18
0
 def run(self):
     parameters.setParameters(parameters.data)
     remote = RemoteContext(SSH_HOST)
     remote.check_output(['mkdir -p ' + parameters.data['remote_storage']])
Example #19
0
 def run(self):
     remote = RemoteContext(SSH_HOST)
     print remote.check_output(["ps aux > {0}".format(self.output().path)])
Example #20
0
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')
Example #21
0
 def setUp(self):
     self.context = RemoteContext(working_ssh_host)