def _download_http_file(self, url, target_path):
'''下载http文件
:param url: HTTP路径
:type url: string
:param target_path: 下载到目标路径
:type target_path: string
'''
url0 = url
if url[:7] == 'http://':
url = url[7:]
pos = url.find('/')
host = url[:pos]
page = url[pos:]
conn = HTTPConnection(host, port=80, timeout=60) # 60秒超时
conn.request('GET', page)
res = conn.getresponse()
if res.status != 200:
raise RuntimeError('访问：%s 错误[HTTP错误码：%s]' % (url0, res.status))
target_size = int(res.getheader('Content-Length'))
data = res.read()
conn.close()
if len(data) != target_size:
return self._download_http_file(url0, target_path)
else:
f = open(target_path, 'wb')
f.write(data)
f.close()
def send_email(request):
try:
recipients = request.GET['to'].split(',')
url = request.GET['url']
proto, server, path, query, frag = urlsplit(url)
if query: path += '?' + query
conn = HTTPConnection(server)
conn.request('GET',path)
try: # Python 2.7+, use buffering of HTTP responses
resp = conn.getresponse(buffering=True)
except TypeError: # Python 2.6 and older
resp = conn.getresponse()
assert resp.status == 200, "Failed HTTP response %s %s" % (resp.status, resp.reason)
rawData = resp.read()
conn.close()
message = MIMEMultipart()
message['Subject'] = "Graphite Image"
message['To'] = ', '.join(recipients)
message['From'] = 'composer@%s' % gethostname()
text = MIMEText( "Image generated by the following graphite URL at %s\r\n\r\n%s" % (ctime(),url) )
image = MIMEImage( rawData )
image.add_header('Content-Disposition', 'attachment', filename="composer_" + strftime("%b%d_%I%M%p.png"))
message.attach(text)
message.attach(image)
s = SMTP(settings.SMTP_SERVER)
s.sendmail('composer@%s' % gethostname(),recipients,message.as_string())
s.quit()
return HttpResponse( "OK" )
except:
return HttpResponse( format_exc() )
def compute_engine_id():
"""Gets the Compute Engine project ID if it can be inferred.
Uses 169.254.169.254 for the metadata server to avoid request
latency from DNS lookup.
See https://cloud.google.com/compute/docs/metadata#metadataserver
for information about this IP address. (This IP is also used for
Amazon EC2 instances, so the metadata flavor is crucial.)
See https://github.com/google/oauth2client/issues/93 for context about
DNS latency.
:rtype: string or ``NoneType``
:returns: Compute Engine project ID if the metadata service is available,
else ``None``.
"""
host = '169.254.169.254'
uri_path = '/computeMetadata/v1/project/project-id'
headers = {'Metadata-Flavor': 'Google'}
connection = HTTPConnection(host, timeout=0.1)
try:
connection.request('GET', uri_path, headers=headers)
response = connection.getresponse()
if response.status == 200:
return response.read()
except socket.error: # socket.timeout or socket.error(64, 'Host is down')
pass
finally:
connection.close()
def send_email(request):
try:
recipients = request.GET['to'].split(',')
url = request.GET['url']
proto, server, path, query, frag = urlsplit(url)
if query: path += '?' + query
conn = HTTPConnection(server)
conn.request('GET',path)
try: # Python 2.7+, use buffering of HTTP responses
resp = conn.getresponse(buffering=True)
except TypeError: # Python 2.6 and older
resp = conn.getresponse()
assert resp.status == 200, "Failed HTTP response %s %s" % (resp.status, resp.reason)
rawData = resp.read()
conn.close()
message = MIMEMultipart()
message['Subject'] = "Graphite Image"
message['To'] = ', '.join(recipients)
message['From'] = 'composer@%s' % gethostname()
text = MIMEText( "Image generated by the following graphite URL at %s\r\n\r\n%s" % (ctime(),url) )
image = MIMEImage( rawData )
image.add_header('Content-Disposition', 'attachment', filename="composer_" + strftime("%b%d_%I%M%p.png"))
message.attach(text)
message.attach(image)
s = SMTP(settings.SMTP_SERVER)
s.sendmail('composer@%s' % gethostname(),recipients,message.as_string())
s.quit()
return HttpResponse( "OK" )
except Exception:
return HttpResponse(format_exc())
def _compute_engine_id():
"""Gets the Compute Engine project ID if it can be inferred.
Uses 169.254.169.254 for the metadata server to avoid request
latency from DNS lookup.
See https://cloud.google.com/compute/docs/metadata#metadataserver
for information about this IP address. (This IP is also used for
Amazon EC2 instances, so the metadata flavor is crucial.)
See https://github.com/google/oauth2client/issues/93 for context about
DNS latency.
:rtype: string or ``NoneType``
:returns: Compute Engine project ID if the metadata service is available,
else ``None``.
"""
host = '169.254.169.254'
uri_path = '/computeMetadata/v1/project/project-id'
headers = {'Metadata-Flavor': 'Google'}
connection = HTTPConnection(host, timeout=0.1)
try:
connection.request('GET', uri_path, headers=headers)
response = connection.getresponse()
if response.status == 200:
return response.read()
except socket.error: # socket.timeout or socket.error(64, 'Host is down')
pass
finally:
connection.close()
class WebcamStreamBase(VideoStream):
request_headers = {
'Accept': '*/*',
}
def __init__(self, url,
max_rate=3.0,
rate_bucket_size=None,
socket_timeout=10,
user_agent=DEFAULT_USER_AGENT):
self.url = url
netloc, self.path = _parse_url(url)
self.conn = HTTPConnection(netloc, timeout=socket_timeout)
self.request_headers = self.request_headers.copy()
self.request_headers['User-Agent'] = user_agent
self.stream = None
self.rate_limiter = BucketRateLimiter(max_rate, rate_bucket_size)
self.open_rate_limiter = BackoffRateLimiter(socket_timeout)
def __repr__(self):
return u"<%s at 0x%x: %s>" % (
self.__class__.__name__, id(self), self.url)
@classmethod
def from_settings(cls, settings, prefix='webcam.', **defaults):
config = config_from_settings(settings, prefix=prefix,
subprefix=cls.settings_subprefix,
**defaults)
return cls(**config)
def close(self):
if self.stream:
self.stream = None
if self.conn:
self.conn.close()
self.conn = None
@property
def closed(self):
return not self.conn
def next(self):
# FIXME: check closed more often?
if self.closed:
raise StopIteration()
next(self.rate_limiter)
try:
if self.stream is None:
next(self.open_rate_limiter)
self.stream = self._open_stream()
frame = next(self.stream)
self.open_rate_limiter.reset()
return frame
except Exception as ex:
self.stream = None
log.warn("Streaming failed: %s", text_type(ex) or repr(ex))
self.conn.close()
return None
def export_to_myexp(self, trans, id, myexp_username, myexp_password):
"""
Exports a workflow to myExperiment website.
"""
trans.workflow_building_mode = workflow_building_modes.ENABLED
stored = self.get_stored_workflow(trans, id, check_ownership=False, check_accessible=True)
# Convert workflow to dict.
workflow_dict = self._workflow_to_dict(trans, stored)
#
# Create and submit workflow myExperiment request.
#
# Create workflow content JSON.
workflow_content = json.dumps(workflow_dict, indent=4, sort_keys=True)
# Create myExperiment request.
request_raw = trans.fill_template(
"workflow/myexp_export.mako",
workflow_name=workflow_dict['name'],
workflow_description=workflow_dict['annotation'],
workflow_content=workflow_content,
workflow_svg=self._workflow_to_svg_canvas(trans, stored).tostring()
)
# strip() b/c myExperiment XML parser doesn't allow white space before XML; utf-8 handles unicode characters.
request = unicodify(request_raw.strip(), 'utf-8')
# Do request and get result.
auth_header = base64.b64encode('%s:%s' % (myexp_username, myexp_password))
headers = {"Content-type": "text/xml", "Accept": "text/xml", "Authorization": "Basic %s" % auth_header}
myexp_url = trans.app.config.myexperiment_target_url
conn = HTTPConnection(myexp_url)
# NOTE: blocks web thread.
conn.request("POST", "/workflow.xml", request, headers)
response = conn.getresponse()
response_data = response.read()
conn.close()
# Do simple parse of response to see if export successful and provide user feedback.
parser = SingleTagContentsParser('id')
parser.feed(response_data)
myexp_workflow_id = parser.tag_content
workflow_list_str = " <br>Return to <a href='%s'>workflow list." % url_for(controller='workflows', action='list')
if myexp_workflow_id:
return trans.show_message(
"""Workflow '%s' successfully exported to myExperiment. <br/>
<a href="http://%s/workflows/%s">Click here to view the workflow on myExperiment</a> %s
""" % (stored.name, myexp_url, myexp_workflow_id, workflow_list_str),
use_panels=True)
else:
return trans.show_error_message(
"Workflow '%s' could not be exported to myExperiment. Error: %s %s" %
(stored.name, response_data, workflow_list_str), use_panels=True)
def export_to_myexp(self, trans, id, myexp_username, myexp_password):
"""
Exports a workflow to myExperiment website.
"""
trans.workflow_building_mode = workflow_building_modes.ENABLED
stored = self.get_stored_workflow(trans, id, check_ownership=False, check_accessible=True)
# Convert workflow to dict.
workflow_dict = self._workflow_to_dict(trans, stored)
#
# Create and submit workflow myExperiment request.
#
# Create workflow content JSON.
workflow_content = json.dumps(workflow_dict, indent=4, sort_keys=True)
# Create myExperiment request.
request_raw = trans.fill_template(
"workflow/myexp_export.mako",
workflow_name=workflow_dict['name'],
workflow_description=workflow_dict['annotation'],
workflow_content=workflow_content,
workflow_svg=self._workflow_to_svg_canvas(trans, stored).tostring()
)
# strip() b/c myExperiment XML parser doesn't allow white space before XML; utf-8 handles unicode characters.
request = unicodify(request_raw.strip(), 'utf-8')
# Do request and get result.
auth_header = base64.b64encode('%s:%s' % (myexp_username, myexp_password))
headers = {"Content-type": "text/xml", "Accept": "text/xml", "Authorization": "Basic %s" % auth_header}
myexp_url = trans.app.config.get("myexperiment_url", self.__myexp_url)
conn = HTTPConnection(myexp_url)
# NOTE: blocks web thread.
conn.request("POST", "/workflow.xml", request, headers)
response = conn.getresponse()
response_data = response.read()
conn.close()
# Do simple parse of response to see if export successful and provide user feedback.
parser = SingleTagContentsParser('id')
parser.feed(response_data)
myexp_workflow_id = parser.tag_content
workflow_list_str = " <br>Return to <a href='%s'>workflow list." % url_for(controller='workflows', action='list')
if myexp_workflow_id:
return trans.show_message(
"""Workflow '%s' successfully exported to myExperiment. <br/>
<a href="http://%s/workflows/%s">Click here to view the workflow on myExperiment</a> %s
""" % (stored.name, myexp_url, myexp_workflow_id, workflow_list_str),
use_panels=True)
else:
return trans.show_error_message(
"Workflow '%s' could not be exported to myExperiment. Error: %s %s" %
(stored.name, response_data, workflow_list_str), use_panels=True)
def is_http_running_on(port):
""" Check if an http server runs on a given port.
Args:
The port to check.
Returns:
True if it is used by an http server. False otherwise.
"""
try:
conn = HTTPConnection('127.0.0.1:' + str(port))
conn.connect()
conn.close()
return True
except Exception:
return False
def _findMatchUrl(self, tag):
h3 = tag.find('h3')
if not h3:
return ''
a = h3.find('a')
url = a.attrs.get('href', '')
# decode url
host = parse.urlsplit(url).netloc
path = url[len(parse.urljoin(url, '/')) - 1:]
conn = HTTPConnection(host, timeout=10)
conn.request('GET', path)
req = conn.getresponse()
r_url = req.getheader('Location')
conn.close()
return r_url
def test_garbage_in(self):
# Connect without SSL regardless of server.scheme
c = HTTPConnection('%s:%s' % (self.interface(), self.PORT))
c._output(b'gjkgjklsgjklsgjkljklsg')
c._send_output()
response = c.response_class(c.sock, method='GET')
try:
response.begin()
self.assertEqual(response.status, 400)
self.assertEqual(response.fp.read(22), b'Malformed Request-Line')
c.close()
except socket.error:
e = sys.exc_info()[1]
# "Connection reset by peer" is also acceptable.
if e.errno != errno.ECONNRESET:
raise
def test_garbage_in(self):
# Connect without SSL regardless of server.scheme
c = HTTPConnection('%s:%s' % (self.interface(), self.PORT))
c._output(b'gjkgjklsgjklsgjkljklsg')
c._send_output()
response = c.response_class(c.sock, method='GET')
try:
response.begin()
self.assertEqual(response.status, 400)
self.assertEqual(response.fp.read(22),
b'Malformed Request-Line')
c.close()
except socket.error:
e = sys.exc_info()[1]
# "Connection reset by peer" is also acceptable.
if e.errno != errno.ECONNRESET:
raise
def _get_version(self, master):
if master is not None:
conn = None
host, port = master.split(':', 2)
port = int(port)
try:
conn = HTTPConnection(host, port, timeout=self._timeout)
conn.request('GET', '/version')
resp = conn.getresponse()
if resp.status < 200 or resp.status >= 300:
return
return json.loads(resp.read().decode('utf-8'))['version']
except Exception:
logger.exception('Error')
pass
finally:
if conn:
conn.close()
def _request(self, method, path, body=None):
payload = None
if body is not None:
try:
payload = json.dumps(body)
except ValueError:
raise ValueError("Failed to encode request body as JSON: {}".format(
json.dumps(body, indent=2)))
if isinstance(payload, text_type):
payload = body.encode("utf-8")
conn = HTTPConnection(self.host, self.port)
try:
conn.request(method, path, payload)
yield conn.getresponse()
finally:
conn.close()
def _get_version(self, master):
if master is not None:
conn = None
host, port = master.split(':', 2)
port = int(port)
try:
conn = HTTPConnection(host, port, timeout=self._timeout)
conn.request('GET', '/version')
resp = conn.getresponse()
if resp.status < 200 or resp.status >= 300:
return
return json.loads(resp.read().decode('utf-8'))['version']
except Exception:
logger.exception('Error')
pass
finally:
if conn:
conn.close()
class GoogleSuggestions():
def __init__(self):
self.hl = "en"
self.conn = None
def prepareQuery(self):
#GET /complete/search?output=toolbar&client=youtube-psuggest&xml=true&ds=yt&hl=en&jsonp=self.gotSuggestions&q=s
#self.prepQuerry = "/complete/search?output=toolbar&client=youtube&xml=true&ds=yt&"
self.prepQuerry = "/complete/search?output=chrome&client=chrome&"
if self.hl is not None:
self.prepQuerry = self.prepQuerry + "hl=" + self.hl + "&"
self.prepQuerry = self.prepQuerry + "jsonp=self.gotSuggestions&q="
print("[MyTube - GoogleSuggestions] prepareQuery:", self.prepQuerry)
def getSuggestions(self, queryString):
self.prepareQuery()
if queryString != "":
query = self.prepQuerry + quote(queryString)
self.conn = HTTPConnection("google.com")
try:
self.conn = HTTPConnection("google.com")
self.conn.request("GET", query, "", {"Accept-Encoding": "UTF-8"})
except (CannotSendRequest, gaierror, error):
self.conn.close()
print("[MyTube - GoogleSuggestions] Can not send request for suggestions")
return None
else:
try:
response = self.conn.getresponse()
except BadStatusLine:
self.conn.close()
print("[MyTube - GoogleSuggestions] Can not get a response from google")
return None
else:
if response.status == 200:
data = response.read()
header = response.getheader("Content-Type", "text/xml; charset=ISO-8859-1")
charset = "ISO-8859-1"
try:
charset = header.split(";")[1].split("=")[1]
print("[MyTube - GoogleSuggestions] Got charset %s" % charset)
except:
print("[MyTube - GoogleSuggestions] No charset in Header, falling back to %s" % charset)
data = data.decode(charset).encode("utf-8")
self.conn.close()
return data
else:
self.conn.close()
return None
else:
return None
class WebcamStreamBase(VideoStream):
request_headers = {
'Accept': '*/*',
}
def __init__(self,
url,
max_rate=3.0,
rate_bucket_size=None,
socket_timeout=10,
user_agent=DEFAULT_USER_AGENT):
self.url = url
netloc, self.path = _parse_url(url)
self.conn = HTTPConnection(netloc, timeout=socket_timeout)
self.request_headers = self.request_headers.copy()
self.request_headers['User-Agent'] = user_agent
self.stream = None
self.rate_limiter = BucketRateLimiter(max_rate, rate_bucket_size)
self.open_rate_limiter = BackoffRateLimiter(socket_timeout)
def __repr__(self):
return u"<%s at 0x%x: %s>" % (self.__class__.__name__, id(self),
self.url)
@classmethod
def from_settings(cls, settings, prefix='webcam.', **defaults):
config = config_from_settings(settings,
prefix=prefix,
subprefix=cls.settings_subprefix,
**defaults)
return cls(**config)
def close(self):
if self.stream:
self.stream = None
if self.conn:
self.conn.close()
self.conn = None
@property
def closed(self):
return not self.conn
def next(self):
# FIXME: check closed more often?
if self.closed:
raise StopIteration()
next(self.rate_limiter)
try:
if self.stream is None:
next(self.open_rate_limiter)
self.stream = self._open_stream()
frame = next(self.stream)
self.open_rate_limiter.reset()
return frame
except Exception as ex:
self.stream = None
log.warn("Streaming failed: %s", text_type(ex) or repr(ex))
self.conn.close()
return None
class MesosSchedulerDriver(Process, SchedulerDriver):
_timeout = 10
def __init__(self,
sched,
framework,
master_uri,
use_addict=False,
implicit_acknowledgements=True,
principal=None,
secret=None,
failover=False,
timeout=DAY):
super(MesosSchedulerDriver, self).__init__(timeout=timeout)
self.sched = sched
self.master_uri = master_uri
self._framework = framework
self.detector = None
self._conn = None
self.version = None
self._failover = failover
self._dict_cls = Dict if use_addict else dict
self.implicit_acknowledgements = implicit_acknowledgements
if principal is not None and secret is not None:
self._basic_credential = 'Basic %s' % (b2a_base64(
('%s:%s' %
(principal, secret)).encode('ascii')).decode('ascii').strip())
else:
self._basic_credential = None
@property
def framework(self):
framework = dict(self._framework)
version = self.version and tuple(
int(n) for n in self.version.split('.'))
capabilities = [
c for c in framework.get('capabilities', [])
if c['type'] != 'GPU_RESOURCES'
]
if version and version >= (1, 0, 0):
capabilities.append(dict(type='GPU_RESOURCES'))
if capabilities:
framework['capabilities'] = capabilities
else:
framework.pop('capabilities', None)
if 'failover_timeout' not in framework:
framework['failover_timeout'] = 100
return framework
@property
def framework_id(self):
id = self._framework.get('id')
return id and id.get('value')
@framework_id.setter
def framework_id(self, id):
self._framework['id'] = dict(value=id)
def _get_version(self, master):
if master is not None:
conn = None
host, port = master.split(':', 2)
port = int(port)
try:
conn = HTTPConnection(host, port, timeout=self._timeout)
conn.request('GET', '/version')
resp = conn.getresponse()
if resp.status < 200 or resp.status >= 300:
return
return json.loads(resp.read().decode('utf-8'))['version']
except Exception:
logger.exception('Error')
pass
finally:
if conn:
conn.close()
def change_master(self, master):
self.version = self._get_version(master)
super(MesosSchedulerDriver, self).change_master(master)
self._close()
def start(self):
super(MesosSchedulerDriver, self).start()
uri = self.master_uri
if uri.startswith('zk://') or uri.startswith('zoo://'):
from .detector import MasterDetector
self.detector = MasterDetector(uri[uri.index('://') + 3:], self)
self.detector.start()
else:
if ':' not in uri:
uri += ':5050'
self.change_master(uri)
def stop(self, failover=False):
with self._lock:
self._failover = failover
detector = self.detector
self.detector = None
if detector:
detector.stop()
super(MesosSchedulerDriver, self).stop()
def _shutdown(self):
if not self._failover:
try:
self._teardown()
except Exception:
logger.exception('Failed to Teardown')
def _get_conn(self):
if not self.connected:
return None
if self._conn is not None:
return self._conn
host, port = self.master.split(':', 2)
port = int(port)
self._conn = HTTPConnection(host, port, timeout=self._timeout)
return self._conn
def _send(self, body, path='/api/v1/scheduler', method='POST', headers={}):
with self._lock:
conn = self._get_conn()
if conn is None:
raise RuntimeError('Not connected yet')
if body != '':
data = json.dumps(body).encode('utf-8')
headers['Content-Type'] = 'application/json'
else:
data = b''
stream_id = self.stream_id
if stream_id:
headers['Mesos-Stream-Id'] = stream_id
if self._basic_credential:
headers['Authorization'] = self._basic_credential
try:
conn.request(method, path, body=data, headers=headers)
resp = conn.getresponse()
except Exception:
self._close()
raise
if resp.status < 200 or resp.status >= 300:
raise RuntimeError('Failed to send request %s: %s\n%s' %
(resp.status, resp.read(), data))
result = resp.read()
if not result:
return {}
try:
return json.loads(result.decode('utf-8'))
except Exception:
return {}
def _teardown(self):
if self.connected:
framework_id = self.framework_id
if framework_id:
self._send(
dict(
type='TEARDOWN',
framework_id=dict(value=framework_id, ),
))
self._framework.pop('id', None)
def acceptOffers(self, offer_ids, operations, filters=None):
if not operations:
return self.declineOffer(offer_ids, filters=filters)
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
accept = dict(
offer_ids=[offer_ids]
if isinstance(offer_ids, dict) else offer_ids,
operations=operations,
)
if filters is not None:
accept['filters'] = filters
body = dict(
type='ACCEPT',
framework_id=dict(value=framework_id, ),
accept=accept,
)
self._send(body)
def acceptInverseOffers(self, offer_ids, filters=None):
framework_id = self.framework_id
assert framework_id
accept_inverse_offers = dict(
inverse_offer_ids=[offer_ids] if isinstance(offer_ids, dict
) else offer_ids)
if filters is not None:
accept_inverse_offers['filters'] = filters
body = dict(
type='ACCEPT_INVERSE_OFFERS',
framework_id=dict(value=framework_id, ),
accept_inverse_offers=accept_inverse_offers,
)
self._send(body)
def launchTasks(self, offer_ids, tasks, filters=None):
if not tasks:
return self.declineOffer(offer_ids, filters=filters)
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
operations = [dict(
type='LAUNCH',
launch=dict(task_infos=tasks),
)]
self.acceptOffers(offer_ids, operations, filters=filters)
def declineOffer(self, offer_ids, filters=None):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
decline = dict(offer_ids=[offer_ids] if isinstance(offer_ids, dict
) else offer_ids)
if filters is not None:
decline['filters'] = filters
body = dict(
type='DECLINE',
framework_id=dict(value=framework_id, ),
decline=decline,
)
self._send(body)
def declineInverseOffer(self, offer_ids, filters=None):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
decline_inverse_offers = dict(
inverse_offer_ids=[offer_ids] if isinstance(offer_ids, dict
) else offer_ids)
if filters is not None:
decline_inverse_offers['filters'] = filters
body = dict(
type='DECLINE_INVERSE_OFFERS',
framework_id=dict(value=framework_id, ),
decline_inverse_offers=decline_inverse_offers,
)
self._send(body)
def reviveOffers(self):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='REVIVE',
framework_id=dict(value=framework_id, ),
)
self._send(body)
def suppressOffers(self):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='SUPPRESS',
framework_id=dict(value=framework_id, ),
)
self._send(body)
def killTask(self, task_id):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='KILL',
framework_id=dict(value=framework_id, ),
kill=dict(task_id=task_id, ),
)
self._send(body)
def acknowledgeStatusUpdate(self, status):
if self.connected and 'uuid' in status:
framework_id = self.framework_id
assert framework_id
acknowledge = dict()
acknowledge['agent_id'] = status['agent_id']
acknowledge['task_id'] = status['task_id']
acknowledge['uuid'] = status['uuid']
body = dict(
type='ACKNOWLEDGE',
framework_id=dict(value=framework_id, ),
acknowledge=acknowledge,
)
self._send(body)
def acknowledgeOperationStatusUpdate(self, status):
if self.connected and 'uuid' in status and 'operation_id' in status:
framework_id = self.framework_id
assert framework_id
body = dict(
type='ACKNOWLEDGE_OPERATION_STATUS',
framework_id=dict(value=framework_id, ),
acknowledge_operation_status=status,
)
self._send(body)
def reconcileTasks(self, tasks):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='RECONCILE',
framework_id=dict(value=framework_id, ),
reconcile=dict(
tasks=[dict(task_id=task['task_id']) for task in tasks], ),
)
self._send(body)
def reconcileOperations(self, operations_):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
operations = []
for op_ in operations_:
op = dict(operation_id=op_['operation_id'])
if 'agent_id' in op_:
op['agent_id'] = op_['agent_id']
if 'resource_provider_id' in op_:
op['resource_provider_id'] = op_['resource_provider_id']
operations.append(op)
body = dict(
type='RECONCILE_OPERATIONS',
framework_id=dict(value=framework_id, ),
reconcile_operations=dict(operations=operations, ),
)
self._send(body)
def sendFrameworkMessage(self, executor_id, agent_id, data):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
message = dict(
agent_id=agent_id,
executor_id=executor_id,
data=data,
)
body = dict(
type='MESSAGE',
framework_id=dict(value=framework_id, ),
message=message,
)
self._send(body)
def requestResources(self, requests):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='REQUEST',
framework_id=dict(value=framework_id, ),
request=dict(requests=requests, ),
)
self._send(body)
def onNewMasterDetectedMessage(self, data):
master = None
try:
if isinstance(data, six.binary_type):
data = data.decode('utf-8')
parsed = json.loads(data)
if parsed and "address" in parsed:
ip = parsed["address"].get("ip")
port = parsed["address"].get("port")
if ip and port:
master = "%s:%s" % (ip, port)
except Exception:
logger.exception("No JSON content, probably connecting "
"to older Mesos version.")
if master:
self.change_master(master)
def onNoMasterDetectedMessage(self):
self.change_master(None)
def gen_request(self):
request = dict(
type='SUBSCRIBE',
subscribe=dict(framework_info=self.framework),
)
if 'id' in self._framework:
request['framework_id'] = self._framework['id']
data = json.dumps(request)
_authorization = ''
if self._basic_credential is not None:
_authorization = 'Authorization: %s\r\n' % (
self._basic_credential, )
request = ('POST /api/v1/scheduler HTTP/1.1\r\nHost: %s\r\n'
'Content-Type: application/json\r\n'
'Accept: application/json\r\n%s'
'Connection: close\r\nContent-Length: %s\r\n\r\n%s') % (
self.master, _authorization, len(data), data)
return request.encode('utf-8')
def _close(self):
if self._conn is not None:
self._conn.close()
self._conn = None
def on_close(self):
self._close()
self.sched.disconnected(self)
def on_subscribed(self, info):
reregistered = (self.framework_id is not None)
self.framework_id = info['framework_id']['value']
hostname, port = self.master.split(':', 2)
port = int(port)
master_info = dict(
hostname=hostname,
port=port,
)
if self.version:
master_info['version'] = self.version
elif 'master_info' in info and 'version' in info['master_info']:
master_info['version'] = info['master_info']['version']
if reregistered:
self.sched.reregistered(self, self._dict_cls(master_info))
else:
framework_id = dict(value=self.framework_id)
self.sched.registered(self, self._dict_cls(framework_id),
self._dict_cls(master_info))
def on_offers(self, event):
offers = event.get('offers', [])
if offers:
self.sched.resourceOffers(
self, [self._dict_cls(offer) for offer in offers])
version = self.version and tuple(
int(n) for n in self.version.split('.'))
if not (version and version >= (1, 0, 0)):
self.on_inverse_offers(event)
def on_inverse_offers(self, event):
inverse_offers = event.get('inverse_offers', [])
if inverse_offers:
self.sched.inverseOffers(
self, [self._dict_cls(offer) for offer in inverse_offers])
def on_rescind(self, event):
offer_id = event['offer_id']
self.sched.offerRescinded(self, self._dict_cls(offer_id))
def on_rescind_inverse_offer(self, event):
inverse_offer_id = event['inverse_offer_id']
self.sched.inverseOfferRescinded(self,
self._dict_cls(inverse_offer_id))
def on_update(self, event):
status = event['status']
self.sched.statusUpdate(self, self._dict_cls(status))
if self.implicit_acknowledgements:
self.acknowledgeStatusUpdate(status)
def on_update_operation_status(self, event):
status = event['status']
self.sched.operationStatusUpdate(self, self._dict_cls(status))
if self.implicit_acknowledgements:
self.acknowledgeOperationStatusUpdate(status)
def on_message(self, message):
executor_id = message['executor_id']
agent_id = message['agent_id']
data = message['data']
self.sched.frameworkMessage(self, self._dict_cls(executor_id),
self._dict_cls(agent_id), data)
def on_failure(self, failure):
agent_id = failure['agent_id']
if 'executor_id' not in failure:
self.sched.slaveLost(self, self._dict_cls(agent_id))
else:
self.sched.executorLost(self,
self._dict_cls(failure['executor_id']),
self._dict_cls(agent_id),
failure['status'])
def on_error(self, event):
message = event['message']
self.sched.error(self, message)
def on_heartbeat(self):
self.sched.processHeartBeat()
def on_event(self, event):
if 'type' in event:
_type = event['type'].lower()
if _type == 'heartbeat':
self.on_heartbeat()
return
if _type not in event:
logger.error('Missing `%s` in event %s' % (_type, event))
return
event = event[_type]
func_name = 'on_%s' % (_type, )
func = getattr(self, func_name, None)
if func is not None:
func(event)
else:
logger.error('Unknown type:%s, event:%s' % (_type, event))
else:
logger.error('Unknown event:%s' % (event, ))
class MesosExecutorDriver(Process, ExecutorDriver):
def __init__(self, executor):
env = os.environ
agent_endpoint = env['MESOS_AGENT_ENDPOINT']
framework_id = env['MESOS_FRAMEWORK_ID']
assert framework_id
self.framework_id = dict(value=framework_id)
executor_id = env['MESOS_EXECUTOR_ID']
self.executor_id = dict(value=executor_id)
grace_shutdown_period = env.get('MESOS_EXECUTOR_SHUTDOWN_GRACE_PERIOD')
if grace_shutdown_period:
self.grace_shutdown_period = parse_duration(grace_shutdown_period)
else:
self.grace_shutdown_period = 0.0
self.checkpoint = bool(env.get('MESOS_CHECKPOINT'))
self.local = bool(env.get('MESOS_LOCAL'))
self.executor = executor
self.framework_info = None
self.executor_info = None
self.tasks = {}
self.updates = {}
self._conn = None
super(MesosExecutorDriver, self).__init__(master=agent_endpoint)
def _delay_kill(self):
def _():
try:
time.sleep(self.grace_shutdown_period)
os.killpg(0, signal.SIGKILL)
except Exception:
logger.exception('Failed to force kill executor')
t = Thread(target=_)
t.daemon = True
t.start()
def gen_request(self):
body = json.dumps(
dict(
type='SUBSCRIBE',
framework_id=self.framework_id,
executor_id=self.executor_id,
subscribe=dict(
unacknowledged_tasks=list(self.tasks.values()),
unacknowledged_updates=list(self.updates.values()),
),
))
request = ('POST /api/v1/executor HTTP/1.1\r\nHost: %s\r\n'
'Content-Type: application/json\r\n'
'Accept: application/json\r\n'
'Connection: close\r\nContent-Length: %s\r\n\r\n%s') % (
self.master, len(body), body)
return request.encode('utf-8')
def on_close(self):
if self._conn is not None:
self._conn.close()
self._conn = None
self.version = None
self.executor.disconnected(self)
if not self.checkpoint:
if not self.local:
self._delay_kill()
self.executor.shutdown(self)
self.abort()
def on_event(self, event):
if 'type' in event:
_type = event['type'].lower()
if _type == 'shutdown':
self.on_shutdown()
return
if _type == 'heartbeat':
return
if _type not in event:
logger.error('Missing `%s` in event %s' % (_type, event))
return
event = event[_type]
func_name = 'on_%s' % (_type, )
func = getattr(self, func_name, None)
if func is not None:
func(event)
else:
logger.error('Unknown type:%s, event:%s' % (_type, event))
else:
logger.error('Unknown event:%s' % (event, ))
def on_subscribed(self, info):
executor_info = info['executor_info']
framework_info = info['framework_info']
agent_info = info['agent_info']
assert executor_info['executor_id'] == self.executor_id
assert framework_info['id'] == self.framework_id
if self.executor_info is None or self.framework_info is None:
self.executor_info = executor_info
self.framework_info = framework_info
self.executor.registered(self, executor_info, framework_info,
agent_info)
else:
self.executor.reregistered(self, agent_info)
def on_launch(self, event):
task_info = event['task']
task_id = task_info['task_id']['value']
assert task_id not in self.tasks
self.tasks[task_id] = task_info
self.executor.launchTask(self, task_info)
def on_kill(self, event):
task_id = event['task_id']
self.executor.killTask(self, task_id)
def on_acknowledged(self, event):
task_id = event['task_id']['value']
uuid_ = uuid.UUID(bytes=a2b_base64(event['uuid']))
self.updates.pop(uuid_, None)
self.tasks.pop(task_id, None)
def on_message(self, event):
data = event['data']
self.executor.frameworkMessage(self, data)
def on_error(self, event):
message = event['message']
self.executor.error(self, message)
def on_shutdown(self):
if not self.local:
self._delay_kill()
self.executor.shutdown(self)
self.abort()
def _get_conn(self):
if not self.connected:
return None
if self._conn is not None:
return self._conn
host, port = self.master.split(':', 2)
port = int(port)
self._conn = HTTPConnection(host, port, timeout=1)
return self._conn
def _send(self, body, path='/api/v1/executor', method='POST', headers={}):
conn = self._get_conn()
if conn is None:
raise RuntimeError('Not connected yet')
if body != '':
data = json.dumps(body).encode('utf-8')
headers['Content-Type'] = 'application/json'
else:
data = ''
stream_id = self.stream_id
if stream_id:
headers['Mesos-Stream-Id'] = stream_id
try:
conn.request(method, path, body=data, headers=headers)
resp = conn.getresponse()
except Exception:
self._conn.close()
self._conn = None
raise
if resp.status < 200 or resp.status >= 300:
raise RuntimeError('Failed to send request code=%s, message=%s' %
(resp.status, resp.read()))
result = resp.read()
if not result:
return {}
try:
return json.loads(result.decode('utf-8'))
except Exception:
return {}
def sendStatusUpdate(self, status):
if 'timestamp' not in status:
status['timestamp'] = int(time.time())
if 'uuid' not in status:
status['uuid'] = b2a_base64(uuid.uuid4().bytes)
if 'source' not in status:
status['source'] = 'SOURCE_EXECUTOR'
body = dict(
type='UPDATE',
executor_id=self.executor_id,
framework_id=self.framework_id,
update=dict(status=status, ),
)
self._send(body)
def sendFrameworkMessage(self, data):
body = dict(
type='MESSAGE',
executor_id=self.executor_id,
framework_id=self.framework_id,
message=dict(data=data, ),
)
self._send(body)
class MesosOperatorDaemonDriver(OperatorDaemonDriver):
_timeout = 10
def __init__(self, daemon_uri):
self.init(daemon_uri)
def init(self, daemon_uri):
"""
:param daemon_uri: masterHost:5050 or agentHost:5051
"""
self._daemon = daemon_uri
self._conn = None
def _get_conn(self):
if self._conn is not None:
return self._conn
host, port = self._daemon.split(':', 2)
port = int(port)
self._conn = HTTPConnection(host, port, timeout=self._timeout)
return self._conn
def _send(self, body, path='/api/v1/operator', method='POST', headers={}):
with self._lock:
conn = self._get_conn()
if conn is None:
raise RuntimeError('Not connected yet')
if body != '':
data = json.dumps(body).encode('utf-8')
headers['Content-Type'] = 'application/json'
else:
data = ''
# stream_id = self.stream_id
# if stream_id:
# headers['Mesos-Stream-Id'] = stream_id
try:
conn.request(method, path, body=data, headers=headers)
resp = conn.getresponse()
except Exception:
self._conn.close()
self._conn = None
raise
if resp.status >= 300 and resp.status <= 399:
url = resp.getheader('location')
parsed = urlparse(url)
self._daemon = '%s:%s' % (parsed.hostname, parsed.port)
self._conn.close()
self._conn = None
return self._send(body, path, method, headers)
if resp.status < 200 or resp.status >= 300:
raise RuntimeError(
'Failed to send request code=%s, message=%s' % (
resp.status, resp.read()
)
)
result = resp.read()
if not result:
return {}
try:
return json.loads(result.decode('utf-8'))
except Exception:
return {}
def getHealth(self):
body = dict(
type='GET_HEALTH',
)
return self._send(body)
def getFlags(self):
body = dict(
type='GET_FLAGS',
)
return self._send(body)
def getVersion(self):
body = dict(
type='GET_VERSION',
)
return self._send(body)
def getMetrics(self, timeout):
body = dict(
type='GET_METRICS',
get_metrics=dict(
timeout=dict(
nanoseconds=timeout,
)
)
)
return self._send(body)
def getLoggingLevel(self):
body = dict(
type='GET_LOGGING_LEVEL',
)
return self._send(body)
def setLoggingLevel(self, level, duration):
body = dict(
type='SET_LOGGING_LEVEL',
set_logging_level=dict(
duration=dict(
nanoseconds=duration,
),
level=level,
)
)
return self._send(body)
def listFiles(self, path):
body = dict(
type='LIST_FILES',
list_files=dict(
path=path,
),
)
return self._send(body)
def readFile(self, path, length, offset):
body = dict(
type='READ_FILE',
read_file=dict(
length=length,
offset=offset,
path=path,
)
)
return self._send(body)
def getState(self):
body = dict(
type='GET_STATE',
)
return self._send(body)
def getFrameworks(self):
body = dict(
type='GET_FRAMEWORKS',
)
return self._send(body)
def getExecutors(self):
body = dict(
type='GET_EXECUTORS',
)
return self._send(body)
def getTasks(self):
body = dict(
type='GET_TASKS',
)
return self._send(body)
class MesosSchedulerDriver(Process, SchedulerDriver):
def __init__(self, sched, framework, master_uri):
super(MesosSchedulerDriver, self).__init__()
self.sched = sched
self.master_uri = master_uri
self._framework = framework
self.detector = None
self._conn = None
self.version = None
@property
def framework(self):
framework = dict(self._framework)
version = self.version and tuple(
int(n) for n in self.version.split('.'))
capabilities = [
c for c in framework.get('capabilities', [])
if c['type'] != 'GPU_RESOURCES'
]
if version and version >= (1, 0, 0):
capabilities.append(dict(type='GPU_RESOURCES'))
if capabilities:
framework['capabilities'] = capabilities
else:
framework.pop('capabilities', None)
return framework
@property
def framework_id(self):
return self._framework.get('id')
@framework_id.setter
def framework_id(self, id):
self._framework['id'] = id
def _get_version(self, master):
if master is not None:
conn = None
host, port = master.split(':', 2)
port = int(port)
try:
conn = HTTPConnection(host, port, timeout=1)
conn.request('GET', '/version')
resp = conn.getresponse()
if resp.status < 200 or resp.status >= 300:
return
return json.loads(resp.read())['version']
except Exception:
logger.exception('Error')
pass
finally:
if conn:
conn.close()
def change_master(self, master):
self.version = self._get_version(master)
super(MesosSchedulerDriver, self).change_master(master)
if self._conn is not None:
self._conn.close()
self._conn = None
def start(self):
super(MesosSchedulerDriver, self).start()
uri = self.master_uri
if uri.startswith('zk://') or uri.startswith('zoo://'):
from .detector import MasterDetector
self.detector = MasterDetector(uri[uri.index('://') + 3:], self)
self.detector.start()
else:
if ':' not in uri:
uri += ':5050'
self.change_master(uri)
def stop(self, failover=False):
if not failover:
try:
self._teardown()
except Exception:
logger.exception('Failed to Teardown')
if self.detector:
self.detector.stop()
super(MesosSchedulerDriver, self).stop()
def _get_conn(self):
if not self.connected:
return None
if self._conn is not None:
return self._conn
host, port = self.master.split(':', 2)
port = int(port)
self._conn = HTTPConnection(host, port, timeout=1)
return self._conn
def _send(self, body, path='/api/v1/scheduler', method='POST', headers={}):
conn = self._get_conn()
if conn is None:
raise RuntimeError('Not connected yet')
if body != '':
data = json.dumps(body).encode('utf-8')
headers['Content-Type'] = 'application/json'
else:
data = ''
stream_id = self.stream_id
if stream_id:
headers['Mesos-Stream-Id'] = stream_id
try:
conn.request(method, path, body=data, headers=headers)
resp = conn.getresponse()
except Exception:
self._conn.close()
self._conn = None
raise
if resp.status < 200 or resp.status >= 300:
raise RuntimeError('Failed to send request %s: %s\n%s' %
(resp.status, resp.read(), data))
result = resp.read()
if not result:
return {}
try:
return json.loads(result.decode('utf-8'))
except Exception:
return {}
def _teardown(self):
framework_id = self.framework_id
if framework_id:
self._send(
dict(
type='TEARDOWN',
framework_id=dict(value=framework_id, ),
))
def acceptOffers(self, offer_ids, operations, filters=None):
if not operations:
return self.declineOffer(offer_ids, filters=filters)
framework_id = self.framework_id
assert framework_id
accept = dict(
offer_ids=offer_ids,
operations=operations,
)
if filters is not None:
accept['filters'] = filters
body = dict(
type='ACCEPT',
framework_id=dict(value=framework_id, ),
accept=accept,
)
self._send(body)
def launchTasks(self, offer_ids, tasks, filters=None):
if not tasks:
return self.declineOffer(offer_ids, filters=filters)
framework_id = self.framework_id
assert framework_id
operations = [dict(
type='LAUNCH',
launch=dict(task_infos=tasks),
)]
self.acceptOffers(offer_ids, operations, filters=filters)
def declineOffer(self, offer_ids, filters=None):
framework_id = self.framework_id
assert framework_id
decline = dict(offer_ids=offer_ids, )
if filters is not None:
decline['filters'] = filters
body = dict(
type='DECLINE',
framework_id=dict(value=framework_id, ),
decline=decline,
)
self._send(body)
def reviveOffers(self):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='REVIVE',
framework_id=dict(value=framework_id, ),
)
self._send(body)
def killTask(self, task_id):
framework_id = self.framework_id
assert framework_id
body = dict(
type='KILL',
framework_id=dict(value=framework_id, ),
kill=dict(task_id=task_id, ),
)
self._send(body)
def acknowledgeStatusUpdate(self, status):
if 'uuid' in status:
framework_id = self.framework_id
assert framework_id
acknowledge = dict()
acknowledge['agent_id'] = status['agent_id']
acknowledge['task_id'] = status['task_id']
acknowledge['uuid'] = status['uuid']
body = dict(
type='ACKNOWLEDGE',
framework_id=dict(value=framework_id, ),
acknowledge=acknowledge,
)
self._send(body)
def reconcileTasks(self, tasks):
framework_id = self.framework_id
assert framework_id
body = dict(
type='RECONCILE',
framework_id=dict(value=framework_id, ),
reconcile=dict(
tasks=[dict(task_id=task['task_id']) for task in tasks], ),
)
self._send(body)
def sendFrameworkMessage(self, executor_id, agent_id, data):
framework_id = self.framework_id
assert framework_id
message = dict(
agent_id=agent_id,
executor_id=executor_id,
data=b2a_base64(data.encode('utf-8')).rstrip(),
)
body = dict(
type='MESSAGE',
framework_id=dict(value=framework_id, ),
message=message,
)
self._send(body)
def requestResources(self, requests):
framework_id = self.framework_id
assert framework_id
body = dict(
type='REQUEST',
framework_id=dict(value=framework_id, ),
request=dict(requests=requests, ),
)
self._send(body)
def onNewMasterDetectedMessage(self, data):
master = None
try:
parsed = json.loads(data)
if parsed and "address" in parsed:
ip = parsed["address"].get("ip")
port = parsed["address"].get("port")
if ip and port:
master = "%s:%s" % (ip, port)
except Exception:
logger.exception("No JSON content, probably connecting "
"to older Mesos version.")
if master:
self.change_master(master)
def onNoMasterDetectedMessage(self):
self.change_master(None)
def gen_request(self):
data = json.dumps(
dict(
type='SUBSCRIBE',
subscribe=dict(framework_info=self.framework),
))
request = ('POST /api/v1/scheduler HTTP/1.1\r\nHost: %s\r\n'
'Content-Type: application/json\r\n'
'Accept: application/json\r\n'
'Connection: close\r\nContent-Length: %s\r\n\r\n%s') % (
self.master, len(data), data)
return request.encode('utf-8')
def on_close(self):
if self._conn is not None:
self._conn.close()
self._conn = None
self.sched.disconnected(self)
def on_subscribed(self, info):
reregistered = (self.framework_id is not None)
self.framework_id = info['framework_id']['value']
hostname, port = self.master.split(':', 2)
port = int(port)
master_info = dict(
hostname=hostname,
port=port,
)
if self.version:
master_info['version'] = self.version
if reregistered:
self.sched.reregistered(self, master_info)
else:
framework_id = dict(value=self.framework_id)
self.sched.registered(self, framework_id, master_info)
def on_offers(self, event):
offers = event['offers']
self.sched.resourceOffers(self, offers)
def on_rescind(self, event):
offer_id = event['offer_id']
self.sched.offerRescinded(self, offer_id)
def on_update(self, event):
status = event['status']
self.sched.statusUpdate(self, status)
self.acknowledgeStatusUpdate(status)
def on_message(self, message):
executor_id = message['executor_id']
agent_id = message['agent_id']
data = message['data']
self.sched.frameworkMessage(self, executor_id, agent_id, data)
def on_failure(self, failure):
agent_id = failure['agent_id']
if 'executor_id' not in failure:
self.sched.slaveLost(self, agent_id)
else:
self.sched.executorLost(self, failure['executor_id'], agent_id,
failure['status'])
def on_error(self, event):
message = event['message']
self.sched.error(self, message)
def on_event(self, event):
if 'type' in event:
_type = event['type'].lower()
if _type == 'heartbeat':
return
if _type not in event:
logger.error('Missing `%s` in event %s' % (_type, event))
return
event = event[_type]
func_name = 'on_%s' % (_type, )
func = getattr(self, func_name, None)
if func is not None:
func(event)
else:
logger.error('Unknown type:%s, event:%s' % (_type, event))
else:
logger.error('Unknown event:%s' % (event, ))
class MesosExecutorDriver(Process, ExecutorDriver):
_timeout = 10
def __init__(self, executor, use_addict=False):
env = os.environ
agent_endpoint = env['MESOS_AGENT_ENDPOINT']
super(MesosExecutorDriver, self).__init__(master=agent_endpoint)
framework_id = env['MESOS_FRAMEWORK_ID']
assert framework_id
self.framework_id = dict(value=framework_id)
executor_id = env['MESOS_EXECUTOR_ID']
self.executor_id = dict(value=executor_id)
grace_shutdown_period = env.get('MESOS_EXECUTOR_SHUTDOWN_GRACE_PERIOD')
if grace_shutdown_period:
self.grace_shutdown_period = parse_duration(grace_shutdown_period)
else:
self.grace_shutdown_period = 0.0
self.checkpoint = bool(env.get('MESOS_CHECKPOINT'))
self.local = bool(env.get('MESOS_LOCAL'))
self.executor = executor
self.framework_info = None
self.executor_info = None
self.tasks = {}
self.updates = {}
self._conn = None
self._dict_cls = Dict if use_addict else dict
def _delay_kill(self):
def _():
try:
time.sleep(self.grace_shutdown_period)
os.killpg(0, signal.SIGKILL)
except Exception:
logger.exception('Failed to force kill executor')
t = Thread(target=_)
t.daemon = True
t.start()
def gen_request(self):
body = json.dumps(dict(
type='SUBSCRIBE',
framework_id=self.framework_id,
executor_id=self.executor_id,
subscribe=dict(
unacknowledged_tasks=list(self.tasks.values()),
unacknowledged_updates=list(self.updates.values()),
),
))
request = ('POST /api/v1/executor HTTP/1.1\r\nHost: %s\r\n'
'Content-Type: application/json\r\n'
'Accept: application/json\r\n'
'Connection: close\r\nContent-Length: %s\r\n\r\n%s') % (
self.master, len(body), body
)
return request.encode('utf-8')
def on_close(self):
if self._conn is not None:
self._conn.close()
self._conn = None
self.version = None
self.executor.disconnected(self)
if not self.checkpoint:
if not self.local:
self._delay_kill()
self.executor.shutdown(self)
self.abort()
def on_event(self, event):
if 'type' in event:
_type = event['type'].lower()
if _type == 'shutdown':
self.on_shutdown()
return
if _type == 'heartbeat':
return
if _type not in event:
logger.error(
'Missing `%s` in event %s' %
(_type, event))
return
event = event[_type]
func_name = 'on_%s' % (_type,)
func = getattr(self, func_name, None)
if func is not None:
func(event)
else:
logger.error('Unknown type:%s, event:%s' % (_type, event))
else:
logger.error('Unknown event:%s' % (event,))
def on_subscribed(self, info):
executor_info = info['executor_info']
framework_info = info['framework_info']
agent_info = info['agent_info']
assert executor_info['executor_id'] == self.executor_id
assert framework_info['id'] == self.framework_id
if self.executor_info is None or self.framework_info is None:
self.executor_info = executor_info
self.framework_info = framework_info
self.executor.registered(
self, self._dict_cls(executor_info),
self._dict_cls(framework_info), self._dict_cls(agent_info)
)
else:
self.executor.reregistered(self, self._dict_cls(agent_info))
def on_launch(self, event):
task_info = event['task']
task_id = task_info['task_id']['value']
assert task_id not in self.tasks
self.tasks[task_id] = task_info
self.executor.launchTask(self, self._dict_cls(task_info))
def on_kill(self, event):
task_id = event['task_id']
self.executor.killTask(self, self._dict_cls(task_id))
def on_acknowledged(self, event):
task_id = event['task_id']['value']
uuid_ = uuid.UUID(bytes=decode_data(event['uuid']))
self.updates.pop(uuid_, None)
self.tasks.pop(task_id, None)
def on_message(self, event):
data = event['data']
self.executor.frameworkMessage(self, data)
def on_error(self, event):
message = event['message']
self.executor.error(self, message)
def on_shutdown(self):
if not self.local:
self._delay_kill()
self.executor.shutdown(self)
self.abort()
def _get_conn(self):
if not self.connected:
return None
if self._conn is not None:
return self._conn
host, port = self.master.split(':', 2)
port = int(port)
self._conn = HTTPConnection(host, port, timeout=self._timeout)
return self._conn
def _send(self, body, path='/api/v1/executor', method='POST', headers={}):
with self._lock:
conn = self._get_conn()
if conn is None:
raise RuntimeError('Not connected yet')
if body != '':
data = json.dumps(body).encode('utf-8')
headers['Content-Type'] = 'application/json'
else:
data = ''
stream_id = self.stream_id
if stream_id:
headers['Mesos-Stream-Id'] = stream_id
try:
conn.request(method, path, body=data, headers=headers)
resp = conn.getresponse()
except Exception:
self._conn.close()
self._conn = None
raise
if resp.status < 200 or resp.status >= 300:
raise RuntimeError(
'Failed to send request code=%s, message=%s' % (
resp.status, resp.read()
)
)
result = resp.read()
if not result:
return {}
try:
return json.loads(result.decode('utf-8'))
except Exception:
return {}
def sendStatusUpdate(self, status):
if 'timestamp' not in status:
status['timestamp'] = int(time.time())
if 'uuid' not in status:
status['uuid'] = encode_data(uuid.uuid4().bytes)
if 'source' not in status:
status['source'] = 'SOURCE_EXECUTOR'
body = dict(
type='UPDATE',
executor_id=self.executor_id,
framework_id=self.framework_id,
update=dict(
status=status,
),
)
self._send(body)
def sendFrameworkMessage(self, data):
body = dict(
type='MESSAGE',
executor_id=self.executor_id,
framework_id=self.framework_id,
message=dict(
data=data,
),
)
self._send(body)
class MesosOperatorDaemonDriver(OperatorDaemonDriver):
_timeout = 10
def __init__(self, daemon_uri):
self.init(daemon_uri)
def init(self, daemon_uri):
"""
:param daemon_uri: masterHost:5050 or agentHost:5051
"""
self._daemon = daemon_uri
self._conn = None
def _get_conn(self):
if self._conn is not None:
return self._conn
host, port = self._daemon.split(':', 2)
port = int(port)
self._conn = HTTPConnection(host, port, timeout=self._timeout)
return self._conn
def _send(self, body, path='/api/v1/operator', method='POST', headers={}):
with self._lock:
conn = self._get_conn()
if conn is None:
raise RuntimeError('Not connected yet')
if body != '':
data = json.dumps(body).encode('utf-8')
headers['Content-Type'] = 'application/json'
else:
data = ''
# stream_id = self.stream_id
# if stream_id:
# headers['Mesos-Stream-Id'] = stream_id
try:
conn.request(method, path, body=data, headers=headers)
resp = conn.getresponse()
except Exception:
self._conn.close()
self._conn = None
raise
if resp.status < 200 or resp.status >= 300:
raise RuntimeError(
'Failed to send request code=%s, message=%s' % (
resp.status, resp.read()
)
)
result = resp.read()
if not result:
return {}
try:
return json.loads(result.decode('utf-8'))
except Exception:
return {}
def getHealth(self):
body = dict(
type='GET_HEALTH',
)
return self._send(body)
def getFlags(self):
body = dict(
type='GET_FLAGS',
)
return self._send(body)
def getVersion(self):
body = dict(
type='GET_VERSION',
)
return self._send(body)
def getMetrics(self, timeout):
body = dict(
type='GET_METRICS',
get_metrics=dict(
timeout=dict(
nanoseconds=timeout,
)
)
)
return self._send(body)
def getLoggingLevel(self):
body = dict(
type='GET_LOGGING_LEVEL',
)
return self._send(body)
def setLoggingLevel(self, level, duration):
body = dict(
type='SET_LOGGING_LEVEL',
set_logging_level=dict(
duration=dict(
nanoseconds=duration,
),
level=level,
)
)
return self._send(body)
def listFiles(self, path):
body = dict(
type='LIST_FILES',
list_files=dict(
path=path,
),
)
return self._send(body)
def readFile(self, path, length, offset):
body = dict(
type='READ_FILE',
read_file=dict(
length=length,
offset=offset,
path=path,
)
)
return self._send(body)
def getState(self):
body = dict(
type='GET_STATE',
)
return self._send(body)
def getFrameworks(self):
body = dict(
type='GET_FRAMEWORKS',
)
return self._send(body)
def getExecutors(self):
body = dict(
type='GET_EXECUTORS',
)
return self._send(body)
def getTasks(self):
body = dict(
type='GET_TASKS',
)
return self._send(body)
class MesosSchedulerDriver(Process, SchedulerDriver):
_timeout = 10
def __init__(self, sched, framework, master_uri,
use_addict=False, implicit_acknowledgements=True,
principal=None, secret=None, failover=False,
timeout=DAY):
super(MesosSchedulerDriver, self).__init__(timeout=timeout)
self.sched = sched
self.master_uri = master_uri
self._framework = framework
self.detector = None
self._conn = None
self.version = None
self._failover = failover
self._dict_cls = Dict if use_addict else dict
self.implicit_acknowledgements = implicit_acknowledgements
if principal is not None and secret is not None:
self._basic_credential = 'Basic %s' % (
b2a_base64(
('%s:%s' % (principal, secret)).encode('ascii')
).decode('ascii').strip()
)
else:
self._basic_credential = None
@property
def framework(self):
framework = dict(self._framework)
version = self.version and tuple(
int(n) for n in self.version.split('.')
)
capabilities = [
c for c in framework.get('capabilities', [])
if c['type'] != 'GPU_RESOURCES'
]
if version and version >= (1, 0, 0):
capabilities.append(dict(type='GPU_RESOURCES'))
if capabilities:
framework['capabilities'] = capabilities
else:
framework.pop('capabilities', None)
if 'failover_timeout' not in framework:
framework['failover_timeout'] = 100
return framework
@property
def framework_id(self):
id = self._framework.get('id')
return id and id.get('value')
@framework_id.setter
def framework_id(self, id):
self._framework['id'] = dict(value=id)
def _get_version(self, master):
if master is not None:
conn = None
host, port = master.split(':', 2)
port = int(port)
try:
conn = HTTPConnection(host, port, timeout=self._timeout)
conn.request('GET', '/version')
resp = conn.getresponse()
if resp.status < 200 or resp.status >= 300:
return
return json.loads(resp.read().decode('utf-8'))['version']
except Exception:
logger.exception('Error')
pass
finally:
if conn:
conn.close()
def change_master(self, master):
self.version = self._get_version(master)
super(MesosSchedulerDriver, self).change_master(master)
self._close()
def start(self):
super(MesosSchedulerDriver, self).start()
uri = self.master_uri
if uri.startswith('zk://') or uri.startswith('zoo://'):
from .detector import MasterDetector
self.detector = MasterDetector(uri[uri.index('://') + 3:], self)
self.detector.start()
else:
if ':' not in uri:
uri += ':5050'
self.change_master(uri)
def stop(self, failover=False):
with self._lock:
self._failover = failover
detector = self.detector
self.detector = None
if detector:
detector.stop()
super(MesosSchedulerDriver, self).stop()
def _shutdown(self):
if not self._failover:
try:
self._teardown()
except Exception:
logger.exception('Failed to Teardown')
def _get_conn(self):
if not self.connected:
return None
if self._conn is not None:
return self._conn
host, port = self.master.split(':', 2)
port = int(port)
self._conn = HTTPConnection(host, port, timeout=self._timeout)
return self._conn
def _send(self, body, path='/api/v1/scheduler', method='POST', headers={}):
with self._lock:
conn = self._get_conn()
if conn is None:
raise RuntimeError('Not connected yet')
if body != '':
data = json.dumps(body).encode('utf-8')
headers['Content-Type'] = 'application/json'
else:
data = b''
stream_id = self.stream_id
if stream_id:
headers['Mesos-Stream-Id'] = stream_id
if self._basic_credential:
headers['Authorization'] = self._basic_credential
try:
conn.request(method, path, body=data, headers=headers)
resp = conn.getresponse()
except Exception:
self._close()
raise
if resp.status < 200 or resp.status >= 300:
raise RuntimeError('Failed to send request %s: %s\n%s' % (
resp.status, resp.read(), data))
result = resp.read()
if not result:
return {}
try:
return json.loads(result.decode('utf-8'))
except Exception:
return {}
def _teardown(self):
if self.connected:
framework_id = self.framework_id
if framework_id:
self._send(dict(
type='TEARDOWN',
framework_id=dict(
value=framework_id,
),
))
self._framework.pop('id', None)
def acceptOffers(self, offer_ids, operations, filters=None):
if not operations:
return self.declineOffer(offer_ids, filters=filters)
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
accept = dict(
offer_ids=[offer_ids]
if isinstance(offer_ids, dict)
else offer_ids,
operations=operations,
)
if filters is not None:
accept['filters'] = filters
body = dict(
type='ACCEPT',
framework_id=dict(
value=framework_id,
),
accept=accept,
)
self._send(body)
def acceptInverseOffers(self, offer_ids, filters=None):
framework_id = self.framework_id
assert framework_id
accept_inverse_offers = dict(
inverse_offer_ids=[offer_ids]
if isinstance(offer_ids, dict)
else offer_ids
)
if filters is not None:
accept_inverse_offers['filters'] = filters
body = dict(
type='ACCEPT_INVERSE_OFFERS',
framework_id=dict(
value=framework_id,
),
accept_inverse_offers=accept_inverse_offers,
)
self._send(body)
def launchTasks(self, offer_ids, tasks, filters=None):
if not tasks:
return self.declineOffer(offer_ids, filters=filters)
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
operations = [dict(
type='LAUNCH',
launch=dict(
task_infos=tasks
),
)]
self.acceptOffers(offer_ids, operations, filters=filters)
def declineOffer(self, offer_ids, filters=None):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
decline = dict(
offer_ids=[offer_ids] if isinstance(offer_ids, dict) else offer_ids
)
if filters is not None:
decline['filters'] = filters
body = dict(
type='DECLINE',
framework_id=dict(
value=framework_id,
),
decline=decline,
)
self._send(body)
def declineInverseOffer(self, offer_ids, filters=None):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
decline_inverse_offers = dict(
inverse_offer_ids=[offer_ids]
if isinstance(offer_ids, dict)
else offer_ids
)
if filters is not None:
decline_inverse_offers['filters'] = filters
body = dict(
type='DECLINE_INVERSE_OFFERS',
framework_id=dict(
value=framework_id,
),
decline_inverse_offers=decline_inverse_offers,
)
self._send(body)
def reviveOffers(self, roles=()):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='REVIVE',
framework_id=dict(
value=framework_id,
),
)
if roles:
body['revive'] = dict(roles=list(roles))
self._send(body)
def suppressOffers(self, roles=()):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='SUPPRESS',
framework_id=dict(
value=framework_id,
),
)
if roles:
body['suppress'] = dict(roles=list(roles))
self._send(body)
def killTask(self, task_id):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='KILL',
framework_id=dict(
value=framework_id,
),
kill=dict(
task_id=task_id,
),
)
self._send(body)
def acknowledgeStatusUpdate(self, status):
if self.connected and 'uuid' in status:
framework_id = self.framework_id
assert framework_id
acknowledge = dict()
acknowledge['agent_id'] = status['agent_id']
acknowledge['task_id'] = status['task_id']
acknowledge['uuid'] = status['uuid']
body = dict(
type='ACKNOWLEDGE',
framework_id=dict(
value=framework_id,
),
acknowledge=acknowledge,
)
self._send(body)
def acknowledgeOperationStatusUpdate(self, status):
if self.connected and 'uuid' in status and 'operation_id' in status:
framework_id = self.framework_id
assert framework_id
body = dict(
type='ACKNOWLEDGE_OPERATION_STATUS',
framework_id=dict(
value=framework_id,
),
acknowledge_operation_status=status,
)
self._send(body)
def reconcileTasks(self, tasks):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='RECONCILE',
framework_id=dict(
value=framework_id,
),
reconcile=dict(
tasks=[dict(task_id=task['task_id']) for task in tasks],
),
)
self._send(body)
def reconcileOperations(self, operations_):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
operations = []
for op_ in operations_:
op = dict(
operation_id=op_['operation_id']
)
if 'agent_id' in op_:
op['agent_id'] = op_['agent_id']
if 'resource_provider_id' in op_:
op['resource_provider_id'] = op_['resource_provider_id']
operations.append(op)
body = dict(
type='RECONCILE_OPERATIONS',
framework_id=dict(
value=framework_id,
),
reconcile_operations=dict(
operations=operations,
),
)
self._send(body)
def sendFrameworkMessage(self, executor_id, agent_id, data):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
message = dict(
agent_id=agent_id,
executor_id=executor_id,
data=data,
)
body = dict(
type='MESSAGE',
framework_id=dict(
value=framework_id,
),
message=message,
)
self._send(body)
def requestResources(self, requests):
if not self.connected:
return
framework_id = self.framework_id
assert framework_id
body = dict(
type='REQUEST',
framework_id=dict(
value=framework_id,
),
request=dict(
requests=requests,
),
)
self._send(body)
def onNewMasterDetectedMessage(self, data):
master = None
try:
if isinstance(data, six.binary_type):
data = data.decode('utf-8')
parsed = json.loads(data)
if parsed and "address" in parsed:
ip = parsed["address"].get("ip")
port = parsed["address"].get("port")
if ip and port:
master = "%s:%s" % (ip, port)
except Exception:
logger.exception("No JSON content, probably connecting "
"to older Mesos version.")
if master:
self.change_master(master)
def onNoMasterDetectedMessage(self):
self.change_master(None)
def gen_request(self):
request = dict(
type='SUBSCRIBE',
subscribe=dict(
framework_info=self.framework
),
)
if 'id' in self._framework:
request['framework_id'] = self._framework['id']
data = json.dumps(request)
_authorization = ''
if self._basic_credential is not None:
_authorization = 'Authorization: %s\r\n' % (
self._basic_credential,
)
request = ('POST /api/v1/scheduler HTTP/1.1\r\nHost: %s\r\n'
'Content-Type: application/json\r\n'
'Accept: application/json\r\n%s'
'Connection: close\r\nContent-Length: %s\r\n\r\n%s') % (
self.master, _authorization, len(data), data
)
return request.encode('utf-8')
def _close(self):
if self._conn is not None:
self._conn.close()
self._conn = None
def on_close(self):
self._close()
self.sched.disconnected(self)
def on_subscribed(self, info):
reregistered = (self.framework_id is not None)
self.framework_id = info['framework_id']['value']
hostname, port = self.master.split(':', 2)
port = int(port)
master_info = dict(
hostname=hostname,
port=port,
)
if self.version:
master_info['version'] = self.version
elif 'master_info' in info and 'version' in info['master_info']:
master_info['version'] = info['master_info']['version']
if reregistered:
self.sched.reregistered(self, self._dict_cls(master_info))
else:
framework_id = dict(
value=self.framework_id
)
self.sched.registered(
self, self._dict_cls(framework_id),
self._dict_cls(master_info)
)
def on_offers(self, event):
offers = event.get('offers', [])
if offers:
self.sched.resourceOffers(
self, [self._dict_cls(offer) for offer in offers]
)
version = self.version and tuple(
int(n) for n in self.version.split('.')
)
if not (version and version >= (1, 0, 0)):
self.on_inverse_offers(event)
def on_inverse_offers(self, event):
inverse_offers = event.get('inverse_offers', [])
if inverse_offers:
self.sched.inverseOffers(
self, [self._dict_cls(offer) for offer in inverse_offers]
)
def on_rescind(self, event):
offer_id = event['offer_id']
self.sched.offerRescinded(self, self._dict_cls(offer_id))
def on_rescind_inverse_offer(self, event):
inverse_offer_id = event['inverse_offer_id']
self.sched.inverseOfferRescinded(
self, self._dict_cls(inverse_offer_id)
)
def on_update(self, event):
status = event['status']
self.sched.statusUpdate(self, self._dict_cls(status))
if self.implicit_acknowledgements:
self.acknowledgeStatusUpdate(status)
def on_update_operation_status(self, event):
status = event['status']
self.sched.operationStatusUpdate(self, self._dict_cls(status))
if self.implicit_acknowledgements:
self.acknowledgeOperationStatusUpdate(status)
def on_message(self, message):
executor_id = message['executor_id']
agent_id = message['agent_id']
data = message['data']
self.sched.frameworkMessage(
self, self._dict_cls(executor_id), self._dict_cls(agent_id), data
)
def on_failure(self, failure):
agent_id = failure['agent_id']
if 'executor_id' not in failure:
self.sched.slaveLost(self, self._dict_cls(agent_id))
else:
self.sched.executorLost(
self, self._dict_cls(failure['executor_id']),
self._dict_cls(agent_id), failure['status']
)
def on_error(self, event):
message = event['message']
self.sched.error(self, message)
def on_heartbeat(self):
self.sched.processHeartBeat()
def on_event(self, event):
if 'type' in event:
_type = event['type'].lower()
if _type == 'heartbeat':
self.on_heartbeat()
return
if _type not in event:
logger.error(
'Missing `%s` in event %s' %
(_type, event))
return
event = event[_type]
func_name = 'on_%s' % (_type,)
func = getattr(self, func_name, None)
if func is not None:
func(event)
else:
logger.error('Unknown type:%s, event:%s' % (_type, event))
else:
logger.error('Unknown event:%s' % (event,))
