def test_include_of_single_file(self):
with open("test1.txt", "w") as f:
f.write("hello world")
filedata = pack_include_files(["test1.txt"])
os.makedirs("outdir")
os.chdir("outdir")
unpack_include_files(filedata)
self.assertEquals(os.listdir("."), ["test1.txt"])
def test_include_of_single_file_with_explicit_location(self):
os.makedirs("indir")
os.makedirs("outdir")
with open("indir/test1.txt", "w") as f:
f.write("hello world")
filedata = pack_include_files(["*.txt"], "./indir")
unpack_include_files(filedata, "./outdir")
self.assertEquals(os.listdir("outdir"), ["test1.txt"])
def test_include_of_single_file_with_explicit_location(self):
os.makedirs("indir")
os.makedirs("outdir")
with open("indir/test1.txt", "w") as f:
f.write("hello world")
filedata = pack_include_files(["*.txt"], "./indir")
unpack_include_files(filedata, "./outdir")
self.assertEquals(os.listdir("outdir"), ["test1.txt"])
def test_include_of_single_file(self):
with open("test1.txt", "w") as f:
f.write("hello world")
filedata = pack_include_files(["test1.txt"])
os.makedirs("outdir")
os.chdir("outdir")
unpack_include_files(filedata)
self.assertEquals(os.listdir("."), ["test1.txt"])
def test_preservation_of_file_mode(self):
with open("test1.sh", "w") as f:
f.write("#!/bin/sh\necho 'hello world'\n")
os.chmod("test1.sh", 0755)
with open("private.txt", "w") as f:
f.write("TOP SECRET DATA\n")
os.chmod("private.txt", 0600)
filedata = pack_include_files(["*.*"])
os.unlink("test1.sh")
os.unlink("private.txt")
unpack_include_files(filedata)
self.assertEquals(os.stat("test1.sh").st_mode & 0777, 0755)
self.assertEquals(os.stat("private.txt").st_mode & 0777, 0600)
def test_preservation_of_file_mode(self):
with open("test1.sh", "w") as f:
f.write("#!/bin/sh\necho 'hello world'\n")
os.chmod("test1.sh", 0755)
with open("private.txt", "w") as f:
f.write("TOP SECRET DATA\n")
os.chmod("private.txt", 0600)
filedata = pack_include_files(["*.*"])
os.unlink("test1.sh")
os.unlink("private.txt")
unpack_include_files(filedata)
self.assertEquals(os.stat("test1.sh").st_mode & 0777, 0755)
self.assertEquals(os.stat("private.txt").st_mode & 0777, 0600)
def test_relative_globbing_and_direcotry_includes(self):
os.makedirs("indir")
os.makedirs("outdir")
os.chdir("indir")
with open("test1.txt", "w") as f:
f.write("hello world")
with open("test2.txt", "w") as f:
f.write("hello world")
os.makedirs("subdir/subsubdir")
os.chdir("subdir/subsubdir")
with open("test3.txt", "w") as f:
f.write("hello world")
os.chdir("../../../outdir")
filedata = pack_include_files(["../indir/*.txt", "../indir/*dir"])
unpack_include_files(filedata)
self.assertEquals(sorted(os.listdir(".")),
["subdir", "test1.txt", "test2.txt"])
self.assertEquals(os.listdir("./subdir"), ["subsubdir"])
self.assertEquals(os.listdir("./subdir/subsubdir"), ["test3.txt"])
def test_relative_globbing_and_direcotry_includes(self):
os.makedirs("indir")
os.makedirs("outdir")
os.chdir("indir")
with open("test1.txt", "w") as f:
f.write("hello world")
with open("test2.txt", "w") as f:
f.write("hello world")
os.makedirs("subdir/subsubdir")
os.chdir("subdir/subsubdir")
with open("test3.txt", "w") as f:
f.write("hello world")
os.chdir("../../../outdir")
filedata = pack_include_files(["../indir/*.txt", "../indir/*dir"])
unpack_include_files(filedata)
self.assertEquals(sorted(os.listdir(".")),
["subdir", "test1.txt", "test2.txt"])
self.assertEquals(os.listdir("./subdir"), ["subsubdir"])
self.assertEquals(os.listdir("./subdir/subsubdir"), ["test3.txt"])
def _prepare_filesystem(self):
test_dir = self.args.get('test_dir')
# in standalone mode we take care of creating
# the files
if test_dir is not None:
if not self.slave:
test_dir = test_dir + '-%d' % os.getpid()
if not os.path.exists(test_dir):
os.makedirs(test_dir)
# Copy over the include files, if any.
# It's inefficient to package them up and then immediately
# unpackage them, but this has the advantage of ensuring
# consistency with how it's done in the distributed case.
includes = self.args.get('include_file', [])
filedata = pack_include_files(includes)
unpack_include_files(filedata, test_dir)
# change to execution directory if asked
logger.debug('chdir %r' % test_dir)
os.chdir(test_dir)
def _prepare_filesystem(self):
test_dir = self.args.get('test_dir')
# in standalone mode we take care of creating
# the files
if test_dir is not None:
if not self.slave:
test_dir = test_dir + '-%d' % os.getpid()
if not os.path.exists(test_dir):
os.makedirs(test_dir)
# Copy over the include files, if any.
# It's inefficient to package them up and then immediately
# unpackage them, but this has the advantage of ensuring
# consistency with how it's done in the distributed case.
includes = self.args.get('include_file', [])
filedata = pack_include_files(includes)
unpack_include_files(filedata, test_dir)
# change to execution directory if asked
logger.debug('chdir %r' % test_dir)
os.chdir(test_dir)
class Client(object):
"""Class to drive a Loads cluster.
Options:
- **frontend**: ZMQ socket to call.
- **timeout**: maximum allowed time for a job to run.
Defaults to 1s.
- **timeout_max_overflow**: maximum timeout overflow allowed.
Defaults to 1.5s
- **timeout_overflows**: number of times in a row the timeout value
can be overflowed per agent. The client keeps a counter of
executions that were longer than the regular timeout but shorter
than **timeout_max_overflow**. When the number goes over
**timeout_overflows**, the usual TimeoutError is raised.
When a agent returns on time, the counter is reset.
- **ssh** ssh tunnel server.
"""
def __init__(self,
frontend=DEFAULT_FRONTEND,
timeout=DEFAULT_TIMEOUT,
timeout_max_overflow=DEFAULT_TIMEOUT_MOVF,
timeout_overflows=DEFAULT_TIMEOUT_OVF,
debug=False,
ctx=None,
ssh=None):
self.ssh = ssh
self.kill_ctx = ctx is None
self.ctx = ctx or zmq.Context()
self.frontend = frontend
self.master = self.ctx.socket(zmq.REQ)
if ssh:
from zmq import ssh
ssh.tunnel_connection(self.master, frontend, self.ssh)
else:
self.master.connect(frontend)
self.poller = zmq.Poller()
self.poller.register(self.master, zmq.POLLIN)
self.timeout = timeout * 1000
self.lock = threading.Lock()
self.timeout_max_overflow = timeout_max_overflow * 1000
self.timeout_overflows = timeout_overflows
self.debug = debug
def execute(self, job, timeout=None, log_exceptions=True):
"""Runs the job
Options:
- **job**: Job to be performed. Can be a :class:`Job`
instance or a string. If it's a string a :class:`Job` instance
will be automatically created out of it.
- **timeout**: maximum allowed time for a job to run.
If not provided, uses the one defined in the constructor.
If the job fails after the timeout, raises a :class:`TimeoutError`.
This method is thread-safe and uses a lock. If you need to execute a
lot of jobs simultaneously on a broker, use the :class:`Pool` class.
"""
if timeout is None:
timeout = self.timeout_max_overflow
try:
duration, res = timed(self.debug)(self._execute)(job, timeout)
except Exception:
# logged, connector replaced.
if log_exceptions:
logger.exception('Failed to execute the job.')
raise
if 'error' in res:
raise ValueError(res['error'])
return res['result']
def close(self):
self.master.setsockopt(zmq.LINGER, 0)
self.master.close()
if self.kill_ctx:
self.ctx.destroy(0)
def _execute(self, job, timeout=None):
if not isinstance(job, Message):
job = Message(**job)
if timeout is None:
timeout = self.timeout_max_overflow
with self.lock:
send(self.master, job.serialize())
while True:
try:
socks = dict(self.poller.poll(timeout))
break
except zmq.ZMQError as e:
if e.errno != errno.EINTR:
raise
if socks.get(self.master) == zmq.POLLIN:
data = recv(self.master)
return json.loads(data)
raise TimeoutError(timeout)
def run(self, args, async=True):
# let's ask the broker how many agents it has
res = self.execute({'command': 'LIST'})
# do we have enough ?
agents = len(res)
agents_needed = args.get('agents', 1)
if len(res) < agents_needed:
msg = 'Not enough agents running on that broker. '
msg += 'Asked: %d, Got: %d' % (agents_needed, agents)
raise ExecutionError(msg)
# let's copy over some files if we need
includes = args.get('include_file', [])
cmd = {
'command': 'CTRL_RUN',
'async': async,
'agents': agents_needed,
'args': args
}
cmd['filedata'] = pack_include_files(includes)
res = self.execute(cmd)
logger.debug('Run on its way')
logger.debug(res)
return res
