class BackPressure:
"""
Wrapper for an iterator to provide
async access with backpressure
"""
def __init__(self, iterator, n):
self.iterator = iter(iterator)
self.back_pressure = BoundedSemaphore(n)
def __aiter__(self):
return self
async def __anext__(self):
await self.back_pressure.acquire()
try:
return next(self.iterator)
except StopIteration:
raise StopAsyncIteration
async def release(self, async_iterator):
"""
release iterator to pipeline the backpressure
"""
async for item in async_iterator:
try:
self.back_pressure.release()
except ValueError:
pass
yield item
class LocalBackend(jetstream.backends.BaseBackend):
def __init__(self, cpus=None, blocking_io_penalty=None):
"""The LocalBackend executes tasks as processes on the local machine.
This contains a semaphore that limits tasks by the number of cpus
that they require. It requires that self.runner be set to get the
event loop, so it's not instantiated until preflight.
:param cpus: If this is None, the number of available CPUs will be
guessed. This cannot be changed after starting the backend.
:param blocking_io_penalty: Delay (in seconds) when a BlockingIOError
prevents a new process from spawning.
:param max_concurrency: Max concurrency limit
"""
super(LocalBackend, self).__init__()
self.cpus = cpus \
or jetstream.settings['backends']['local']['cpus'].get() \
or jetstream.utils.guess_local_cpus()
self.bip = blocking_io_penalty \
or jetstream.settings['backends']['local']['blocking_io_penalty'].get(int)
self._cpu_sem = BoundedSemaphore(self.cpus)
log.info(f'LocalBackend initialized with {self.cpus} cpus')
async def spawn(self, task):
log.debug('Spawn: {}'.format(task))
if 'cmd' not in task.directives:
return task.complete()
cmd = task.directives['cmd']
cpus = task.directives.get('cpus', 0)
cpus_reserved = 0
open_fps = list()
if cpus > self.cpus:
raise RuntimeError('Task cpus greater than available cpus')
try:
for i in range(task.directives.get('cpus', 0)):
await self._cpu_sem.acquire()
cpus_reserved += 1
stdin, stdout, stderr = self.get_fd_paths(task)
if stdin:
stdin_fp = open(stdin, 'r')
open_fps.append(stdin_fp)
else:
stdin_fp = None
if stdout:
stdout_fp = open(stdout, 'w')
open_fps.append(stdout_fp)
else:
stdout_fp = None
if stderr:
stderr_fp = open(stderr, 'w')
open_fps.append(stderr_fp)
else:
stderr_fp = None
p = await self.subprocess_sh(cmd,
stdin=stdin_fp,
stdout=stdout_fp,
stderr=stderr_fp)
task.state.update(
stdout_path=stdout,
stderr_path=stderr,
label=f'Slurm({p.pid})',
)
log.info(f'LocalBackend spawned({p.pid}): {task.name}')
rc = await p.wait()
if rc != 0:
log.info(f'Failed: {task.name}')
return task.fail(p.returncode)
else:
log.info(f'Complete: {task.name}')
return task.complete(p.returncode)
except CancelledError:
task.state['err'] = 'Runner cancelled Backend.spawn()'
return task.fail(-15)
finally:
for fp in open_fps:
fp.close()
for i in range(cpus_reserved):
self._cpu_sem.release()
return task
async def subprocess_sh(self,
args,
*,
stdin=None,
stdout=None,
stderr=None,
cwd=None,
encoding=None,
errors=None,
env=None,
loop=None,
executable='/bin/bash'):
"""Asynchronous version of subprocess.run
This will always use a shell to launch the subprocess, and it prefers
/bin/bash (can be changed via arguments)"""
log.debug(f'subprocess_sh:\n{args}')
while 1:
try:
p = await create_subprocess_shell(args,
stdin=stdin,
stdout=stdout,
stderr=stderr,
cwd=cwd,
encoding=encoding,
errors=errors,
env=env,
loop=loop,
executable=executable)
break
except BlockingIOError as e:
log.warning(f'System refusing new processes: {e}')
await asyncio.sleep(self.bip)
return p
async def _check_ports(target: str,
port: int,
loop: asyncio.AbstractEventLoop,
sem: asyncio.BoundedSemaphore,
results: list,
config: DockerScanModel):
open_ports = set()
# for port in ports:
log.error(" > Trying {}:{}".format(target, port))
is_ssl = True
try:
# If connection SSL?
try:
# This definition of ssl context allow to connect with servers with
# self-signed certs
sslcontext = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
sslcontext.options |= ssl.OP_NO_SSLv2
sslcontext.options |= ssl.OP_NO_SSLv3
sslcontext.options |= getattr(ssl, "OP_NO_COMPRESSION", 0)
sslcontext.set_default_verify_paths()
reader, writer = await _get_connection(target,
port,
sslcontext,
config.timeout,
loop)
if not reader:
return
except ssl.SSLError:
reader, writer = await _get_connection(target,
port,
None,
config.timeout,
loop)
if not reader:
return
is_ssl = False
# Send HTTP Header
writer.write(
"GET /v2/ HTTP/1.1\r\nHost: {}\r\n\r\n".format(target).encode()
)
# Get Server response
reader = reader.read(1000)
try:
data = await asyncio.wait_for(reader,
1,
loop=loop)
except (asyncio.TimeoutError, ConnectionRefusedError):
# If this point reached -> server doesn't sent response
return
if b"registry/2.0" in data or \
b"Docker-Distribution-Api-Version" in data:
content = data.lower()
if b"200 ok" in content:
status = "open"
elif b"401" in content:
status = "auth required"
else:
status = "reachable"
log.info(" + Discovered port {}:{}".format(
target,
port
))
open_ports.add((port, status, is_ssl))
# close descriptor
writer.close()
if open_ports:
results.append(
{
target: open_ports
}
)
finally:
sem.release()
class CloudSwiftBackend(BaseBackend):
"""
Executes tasks on cloud-based worker nodes, handling all data transfer to/from worker nodes and
the client node.
"""
def __init__(self,
pw_pool_name=None,
pw_api_key=None,
cpus=None,
blocking_io_penalty=None,
cloud_storage_provider='azure',
**kwargs):
"""
If ``pw_pool_name`` and ``pw_api_key`` are both given valid values, this backend will use ParallelWorks to
manage resouces elastically. If they remain ``None``, then it is assumed that the use has manually started
a pool using the included ``start_pool.py`` script. Note that both approaches requires the use of binaries
available as part of the Swift workflow language.
:param pw_pool_name: str Name of an active pool in PW
:param pw_api_key: str Valid API for a PW account
:param cpus: int The total number of CPUs available to the worker pool
:param blocking_io_penalty: int Delay (in seconds) when a BlockingIOError prevents a new process from spawning.
:param cloud_storage_provder: str Name of the cloud storage provider, which must match up with one of the keys
in ``jetstream.cloud.base.CLOUD_STORAGE_PROVIDERS``
"""
super().__init__()
self.is_pw_pool = pw_pool_name is not None
if pw_pool_name is not None and pw_api_key is not None:
self.pool_info = get_pool_info(pw_pool_name, pw_api_key)
log.info('PW Pool info: {}'.format(self.pool_info))
else:
self.pool_info = {
'cpus': kwargs['pool_info']['cpus_per_worker'],
'maxworkers': kwargs['pool_info']['workers'],
'serviceurl': kwargs['pool_info']['serviceurl'],
}
self.total_cpus_in_pool = self.pool_info['cpus'] * self.pool_info[
'maxworkers']
self.bip = blocking_io_penalty \
or jetstream.settings['backends']['local']['blocking_io_penalty'].get(int)
self._cpu_sem = BoundedSemaphore(int(self.total_cpus_in_pool))
self.project_dir = os.getcwd()
try:
os.remove('cjs_cmds_debug.sh')
except:
pass # Fail silently
# Make directory for cjs launch scripts
self.cloud_scripts_dir = os.path.join(self.project_dir,
'cloud_scripts')
os.makedirs(self.cloud_scripts_dir, exist_ok=True)
self.cloud_logs_dir = os.path.join(self.project_dir, 'cloud_logs')
os.makedirs(self.cloud_logs_dir, exist_ok=True)
# Instantiate a cloud storage provider
storage_class = dynamic_import(
CLOUD_STORAGE_PROVIDERS[cloud_storage_provider])
self.cloud_storage = storage_class(
**kwargs.get(storage_class.config_key, dict()))
# Initialize cloud metrics log
CloudMetricsLogger.init(at=os.path.join(
self.project_dir, 'cloud_metrics_{}.yaml'.format(
datetime.now().strftime('%Y%m%d%H%M%S'))))
# Get path to Petalink library, if it exists
self.petalink_so_path = kwargs.get('petalink_path',
'/usr/lib/petalink.so')
log.info(
f'CloudSwiftBackend initialized with {self.total_cpus_in_pool} cpus'
)
async def spawn(self, task):
# Ensure the command body exists, otherwise there is nothing to do
if 'cmd' not in task.directives:
return task.complete()
# Ensure there will ever be enough CPUs to run this task, otherwise fail
task_requested_cpus = task.directives.get('cpus', 0)
if task_requested_cpus > self.total_cpus_in_pool:
log.critical(
'Task requested cpus ({}) greater than total available cpus ({})'
.format(task_requested_cpus, self.total_cpus_in_pool))
return task.fail(1)
# Add in the pre- and post-hooks into the task body
task.directives['cmd'] = (
f'if [[ -f "{self.petalink_so_path}" ]]; ' +
f'then export LD_PRELOAD={self.petalink_so_path}; fi;\n\n' +
self.cloud_storage.task_cmd_prehook() + '\n' +
task.directives['cmd'] + '\n' +
self.cloud_storage.task_cmd_posthook())
# Determine whether this task should be run locally or on a remote cloud worker
is_local_task = task.directives.get('cloud_args',
dict()).get('local_task', False)
log.info('Spawn ({}): {}'.format('Local' if is_local_task else 'Cloud',
task))
if is_local_task:
return await self.spawn_local(task)
# This is a cloud task
return await self.spawn_cloud(task)
async def spawn_cloud(self, task):
cmd = task.directives['cmd']
cpus_reserved = 0
start_time = datetime.now()
bytes_sent_bundle, bytes_received_bundle = list(), list()
try:
# Get file descriptor paths and file pointers for this task
fd_paths = {
fd_name: fd
for fd_name, fd in zip(('stdin', 'stdout',
'stderr'), self.get_fd_paths(task))
}
fd_filepointers = {
fd_name: open(fd, fd_mode) if fd else None
for fd_mode, (fd_name,
fd) in zip(('r', 'w', 'w'), fd_paths.items())
}
# Upload data inputs into cloud storage
data_metrics = blob_inputs_to_remote(task.directives['input'],
self.cloud_storage)
# Upload reference inputs into cloud storage
reference_inputs = parse_reference_input(
task.directives.get('cloud_args',
dict()).get('reference_input', list()))
ref_metrics = blob_inputs_to_remote(reference_inputs,
self.cloud_storage,
blob_basename=True)
# If the user provides a non-URL path to a container and explicitly says it should be transfer,
# then consider it similar to reference data and upload it to cloud storage
singularity_container_uri = task.directives.get(
'cloud_args', dict()).get('singularity_container')
container_input = ([singularity_container_uri] if (
singularity_container_uri is not None
and not urllib.parse.urlparse(singularity_container_uri).scheme
and get_cloud_directive('transfer_container_to_remote',
task.directives)) else list())
container_metrics = blob_inputs_to_remote(container_input,
self.cloud_storage)
# Log metrics for input data
total_input_metrics = data_metrics + ref_metrics + container_metrics
for m in total_input_metrics:
bytes_sent_bundle.append(m)
bytes_sent_bundle.append({
'name':
'total',
'size':
sum([max(0, t['size']) for t in total_input_metrics]),
'time':
sum([max(0, t['time']) for t in total_input_metrics])
})
# Construct the cog-job-submit command for execution
cjs_cmd = construct_cjs_cmd(
task_body=cmd,
service_url='http://beta.parallel.works:{}'.format(
self.pool_info['serviceport'])
if self.is_pw_pool else self.pool_info['serviceurl'],
cloud_storage=self.cloud_storage,
cjs_stagein=None,
cjs_stageout=None,
cloud_downloads=task.directives['input'] + reference_inputs +
container_input,
cloud_uploads=task.directives['output'],
cloud_scripts_dir=self.cloud_scripts_dir,
singularity_container_uri=singularity_container_uri,
task_name=task.name)
log.debug(cjs_cmd)
# Async submit as a subprocess
p = await self.subprocess_sh(cjs_cmd, **fd_filepointers)
# Once command is executed, update task with some process metadata
task.state.update(
stdout_path=fd_paths['stdout'],
stderr_path=fd_paths['stderr'],
label=f'CloudSwift({p.pid})',
)
log.info(f'CloudSwiftBackend spawned({p.pid}): {task.name}')
rc = await p.wait()
if rc != 0:
log.info(f'Failed: {task.name}')
return task.fail(p.returncode)
else:
# Download completed data files from cloud storage
output_metrics = blob_outputs_to_local(
task.directives['output'], self.cloud_storage)
log.info(f'Complete: {task.name}')
# Log metrics for output data
for m in output_metrics:
bytes_received_bundle.append(m)
bytes_received_bundle.append({
'name':
'total',
'size':
sum([max(0, t['size']) for t in output_metrics]),
'time':
sum([max(0, t['time']) for t in output_metrics])
})
return task.complete(p.returncode)
except CancelledError:
task.state['err'] = 'Runner cancelled Backend.spawn()'
return task.fail(-15)
except Exception as e:
log.error('Exception: {}'.format(e))
traceback.print_exc()
raise
finally:
for fp in fd_filepointers.values():
if fp is not None:
fp.close()
for i in range(cpus_reserved):
self._cpu_sem.release()
# Get task runtime and which node it ran on
elapsed_time = datetime.now() - start_time
hostname = 'UNAVAILABLE'
try:
with open(f'.{task.name}.hostname', 'r') as hostname_log:
hostname = hostname_log.read().strip()
subprocess.call(['mv'] + glob.glob('*.remote.out') +
[self.cloud_logs_dir])
subprocess.call(['mv'] + glob.glob('*.remote.err') +
[self.cloud_logs_dir])
os.remove(f'.{task.name}.hostname')
except:
pass # Fail silently
CloudMetricsLogger.write_record({
'task':
task.name,
'start_datetime':
start_time.strftime('%Y-%m-%d %H:%M:%S'),
'elapsed_time':
str(elapsed_time),
'end_datetime':
datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
'in_files':
bytes_sent_bundle,
'out_files':
bytes_received_bundle,
'node':
hostname
})
return task
async def spawn_local(self, task):
cpus_reserved = 0
fd_filepointers = dict()
try:
for i in range(task.directives.get('cpus', 0)):
await self._cpu_sem.acquire()
cpus_reserved += 1
fd_paths = {
fd_name: fd
for fd_name, fd in zip(('stdin', 'stdout',
'stderr'), self.get_fd_paths(task))
}
fd_filepointers = {
fd_name: open(fd, fd_mode) if fd else None
for fd_mode, (fd_name,
fd) in zip(('r', 'w', 'w'), fd_paths.items())
}
p = await self.subprocess_sh(task.directives['cmd'],
**fd_filepointers)
task.state.update(
stdout_path=fd_paths['stdout'],
stderr_path=fd_paths['stderr'],
label=f'Slurm({p.pid})',
)
log.info(f'LocalBackend spawned({p.pid}): {task.name}')
rc = await p.wait()
if rc != 0:
log.info(f'Failed: {task.name}')
return task.fail(p.returncode)
else:
log.info(f'Complete: {task.name}')
return task.complete(p.returncode)
except CancelledError:
task.state['err'] = 'Runner cancelled Backend.spawn()'
return task.fail(-15)
finally:
for fp in fd_filepointers.values():
if fp is not None:
fp.close()
for i in range(cpus_reserved):
self._cpu_sem.release()
return task
async def subprocess_sh(self,
args,
*,
stdin=None,
stdout=None,
stderr=None,
cwd=None,
encoding=None,
errors=None,
env=None,
loop=None,
executable="/bin/bash"):
"""Asynchronous version of subprocess.run
This will always use a shell to launch the subprocess, and it prefers
/bin/bash (can be changed via arguments)"""
log.debug(f'subprocess_sh:\n{args}')
while 1:
try:
p = await create_subprocess_shell(args,
stdin=stdin,
stdout=stdout,
stderr=stderr,
cwd=cwd,
encoding=encoding,
errors=errors,
env=env,
loop=loop,
executable=executable)
break
except BlockingIOError as e:
log.warning(f'System refusing new processes: {e}')
await asyncio.sleep(self.bip)
return p
async def main():
semaphore = BoundedSemaphore(1)
await semaphore.acquire()
semaphore.release()
semaphore.release()
class HCIHost:
def __init__(self, name: str, dev: str, mon: Monitor = None, **kwargs):
self._dev = dev
self._kwargs = kwargs
self._mon = mon
if (name == HCI_TRANSPORT_UART):
self._transport = UART()
elif (name == HCI_TRANSPORT_TCP):
self._transport = UARToTCP()
else:
raise RuntimeError('Unknown transport type {}'.name)
self._tx_cmd_q = Queue()
# Bound to max 1, so we send only one command at a time
self._tx_cmd_sem = BoundedSemaphore(value=1)
async def open(self):
try:
await self._transport.open(self._dev, **self._kwargs)
except OSError as e:
#log.error(f'Unable to open serial port {e}')
raise e from None
# Start RX task
self._rx_task = create_task(self._rx_task())
# Start TX command task
self._curr_cmd = None
self._tx_cmd_task = create_task(self._tx_cmd_task())
await self.init()
async def close(self):
await self._rx_task
await self._tx_cmd_task
self._transport.close()
async def init(self):
# Reset the controller first
log.debug('init: Reset')
pkt = HCICmd(cmd.Reset)
await self.send_cmd(pkt)
log.debug('init: ReadLocalVersionInformation')
pkt = HCICmd(cmd.ReadLocalVersionInformation)
await self.send_cmd(pkt)
async def _rx_task(self):
log.debug('rx task started')
while True:
try:
pkt = await self._transport.recv()
except CancelledError:
# Here for compatibility with 3.7
raise
except Exception as e:
log.debug(f'_rx_task: {e}')
else:
if not pkt:
# Broken RX Path
log.debug('_rx_task exiting')
break
elif self._mon:
self._mon.feed_rx(0, pkt)
if isinstance(pkt, HCIEvt):
self._rx_evt(pkt)
elif instance(pkt, HCIACLData):
self._rx_acl(pkt)
else:
log.error('Invalid rx type: {type(pkt)}')
def _rx_evt(self, evt: HCIEvt):
log.debug(f'evt rx: {evt}')
# Look for a handler
try:
handler = self.evt_handlers[evt.hdr.code]
except KeyError:
log.warn(f'Discarding event with code: {evt.hdr.code}')
else:
handler(self, evt)
def _rx_acl(self, acl: HCIACLData):
log.debug(f'acl rx: {evt}')
async def _tx_cmd_task(self):
log.debug('tx cmd task started')
while True:
pkt = await self._tx_cmd_q.get()
log.debug(f'pkt tx: {pkt}')
# Wait for the current command to complete
try:
await wait_for(self._tx_cmd_sem.acquire(), 10)
except CancelledError:
# Here for compatibility with 3.7
raise
except TimeoutError:
log.debug('_tx_cmd_task: sem timeout: exiting')
return
assert self._curr_cmd == None
self._curr_cmd = pkt
log.debug(f'_tx_cmd curr set')
try:
await self._transport.send(pkt)
except OSError as e:
log.debug('_tx_cmd_task: send error: exiting')
return
else:
if self._mon:
self._mon.feed_tx(0, pkt)
def tx_cmd(self, pkt: HCICmd) -> None:
self._tx_cmd_q.put_nowait(pkt)
async def send_cmd(self, pkt: HCICmd) -> HCIEvt:
pkt.event.clear()
self.tx_cmd(pkt)
try:
await wait_for(pkt.event.wait(), 15)
except TimeoutError:
raise
def _complete_cmd(self):
log.debug(f'complete: {self._curr_cmd}')
assert self._curr_cmd
# Wake up a potential task waiting on completion
self._curr_cmd.event.set()
self._curr_cmd = None
# The command itself is complete, allow the tx cmd task to move on to
# the next queued command
self._tx_cmd_sem.release()
@evt_handler(evt.CommandComplete)
def _evt_cc(self, evt: HCIEvt) -> None:
log.debug(f'handling CC: {evt}')
self._complete_cmd()
async def _check_ports(target: str, port: int, loop: asyncio.AbstractEventLoop,
sem: asyncio.BoundedSemaphore, results: list,
config: DockerScanModel):
open_ports = set()
# for port in ports:
log.error(" > Trying {}:{}".format(target, port))
is_ssl = True
try:
# If connection SSL?
try:
# This definition of ssl context allow to connect with servers with
# self-signed certs
sslcontext = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
sslcontext.options |= ssl.OP_NO_SSLv2
sslcontext.options |= ssl.OP_NO_SSLv3
sslcontext.options |= getattr(ssl, "OP_NO_COMPRESSION", 0)
sslcontext.set_default_verify_paths()
reader, writer = await _get_connection(target, port, sslcontext,
config.timeout, loop)
if not reader:
return
except ssl.SSLError:
reader, writer = await _get_connection(target, port, None,
config.timeout, loop)
if not reader:
return
is_ssl = False
# Send HTTP Header
writer.write(
"GET /v2/ HTTP/1.1\r\nHost: {}\r\n\r\n".format(target).encode())
# Get Server response
reader = reader.read(1000)
try:
data = await asyncio.wait_for(reader, 1, loop=loop)
except (asyncio.TimeoutError, ConnectionRefusedError):
# If this point reached -> server doesn't sent response
return
if b"registry/2.0" in data or \
b"Docker-Distribution-Api-Version" in data:
content = data.lower()
if b"200 ok" in content:
status = "open"
elif b"401" in content:
status = "auth required"
else:
status = "reachable"
log.info(" + Discovered port {}:{}".format(target, port))
open_ports.add((port, status, is_ssl))
# close descriptor
writer.close()
if open_ports:
results.append({target: open_ports})
finally:
sem.release()