def logs(self,
job_name: str,
task_role: str = 'main',
index: int = 0,
log_type: str = 'stdout'):
assert log_type in [
"stdout", "stderr"
], "now only support stdout and stderr, not %s" % log_type
try:
self.wait([job_name],
exit_states=states_completed + states_unfinished)
info = self.jobs(job_name)
container = info['taskRoles'][task_role]['taskStatuses'][index]
container_id = container['containerId']
if info['jobStatus']['state'] in states_completed:
path_fmt = "http://{ip}:8188/ws/v1/applicationhistory/containers/{container_id}/logs/{stream}?redirected_from_node=true"
ip = re.search('://([\d.]+)/yarn',
container['containerLog']).group(1)
else:
ip = container['containerIp']
path_fmt = "http://{ip}:8042/ws/v1/node/containers/{container_id}/logs/{stream}"
path = path_fmt.format(ip=ip,
container_id=container_id,
stream=log_type)
return get_response(path, method="GET").content
except:
return None
def single_task_logs(status: dict,
task_role: str = 'main',
index: int = 0,
log_type: dict = None,
return_urls: bool = False):
"""change to use containerLog"""
log_type = na(
log_type, {
"stdout": "user.pai.stdout/?start=0",
"stderr": "user.pai.stderr/?start=0"
})
containers = status.get("taskRoles",
{}).get(task_role, {}).get("taskStatuses", [])
if len(containers) < index + 1:
return None
containerLog = containers[index].get("containerLog", None)
if not containerLog:
return None
urls = {k: "{}{}".format(containerLog, v) for k, v in log_type.items()}
if return_urls:
return urls
else:
html_contents = {
k: get_response('GET', v).text
for k, v in urls.items()
}
try:
from html2text import html2text
return {k: html2text(v) for k, v in html_contents.items()}
except ImportError:
return html_contents
def rest_api_job_list(self, user: str = None):
if user is None:
pth = '{}/rest-server/api/v1/jobs'.format(self.pai_uri)
else:
pth = '{}/rest-server/api/v1/user/{}/jobs'.format(
self.pai_uri, user)
return get_response(pth, headers={}, method='GET').json()
def rest_api_token(self, expiration=3600):
return get_response('POST', [self.rest_srv, 'v1', 'token'],
body={
'username': self.user,
'password': self.password,
'expiration': expiration
}).json()['token']
def rest_api_token(self, expiration=3600):
return get_response('{}/rest-server/api/v1/token'.format(self.pai_uri),
body={
'username': self.user,
'password': self.password,
'expiration': expiration
}).json()['token']
def rest_api_token(self, expiration=3600):
__logger__.debug("getting token @ time %s", time.time())
return get_response('{}/rest-server/api/v1/token'.format(self.pai_uri),
body={
'username': self.user,
'password': self.password,
'expiration': expiration
}).json()['token']
def rest_api_user(self, user: str = None):
return get_response(
'GET',
[self.rest_srv, 'v1', 'user', user if user else self.user],
headers={
'Authorization': 'Bearer {}'.format(self.token),
},
).json()
def rest_api_virtual_clusters(self):
return get_response('GET', [self.rest_srv, 'v1', 'virtual-clusters'],
headers={
'Authorization':
'Bearer {}'.format(self.token),
'Content-Type': 'application/json',
},
allowed_status=[200]).json()
def rest_api_jobs(self, job_name: str=None, info: str=None, user: str=None):
user = self.user if user is None else user
pth = '{}/rest-server/api/v1/user/{}/jobs'.format(self.pai_uri, user)
if job_name:
pth = pth + '/' + job_name
if info:
assert info in ['config', 'ssh'], ('unsupported query information', info)
pth = pth + '/' + info
return get_response(pth, headers = {}, method='GET').json()
def rest_api_user(self):
return get_response(
'{}/rest-server/api/v1/user/{}'.format(self.pai_uri, self.user),
method='GET',
headers={
'Authorization': 'Bearer {}'.format(self.token),
'Content-Type': 'text/yaml',
},
).json()
def rest_api_virtual_clusters(self):
return get_response(
'{}/rest-server/api/v1/virtual-clusters'.format(self.pai_uri),
headers={
'Authorization': 'Bearer {}'.format(self.token),
'Content-Type': 'application/json',
},
method='GET',
allowed_status=[200]).json()
def rest_api_submit(self, job: dict):
use_v2 = str(job.get("protocolVersion", 1)) == "2"
if use_v2:
return get_response(
'{}/rest-server/api/v2/jobs'.format(self.pai_uri),
headers={
'Authorization': 'Bearer {}'.format(self.token),
'Content-Type': 'text/yaml',
},
body=job,
allowed_status=[202, 201])
else:
return get_response('{}/rest-server/api/v1/user/{}/jobs'.format(
self.pai_uri, self.user),
headers={
'Authorization':
'Bearer {}'.format(self.token),
'Content-Type': 'application/json',
},
body=job,
allowed_status=[202, 201])
def rest_api_submit(self, job: dict):
use_v2 = str(job.get("protocolVersion", 1)) == "2"
if use_v2:
import yaml
return get_response('POST', [self.rest_srv, 'v2', 'jobs'],
headers={
'Authorization':
'Bearer {}'.format(self.token),
'Content-Type': 'text/yaml',
},
body=yaml.dump(job),
allowed_status=[202, 201])
else:
return get_response(
'POST', [self.rest_srv, 'v1', 'user', self.user, 'jobs'],
headers={
'Authorization': 'Bearer {}'.format(self.token),
'Content-Type': 'application/json',
},
body=job,
allowed_status=[202, 201])
def logs(self,
status: dict = None,
task_role: str = 'main',
index: int = 0,
logs: dict = None):
status = na(status, self.status())
urls = self.get_logs_url(status, task_role, index, logs)
if not urls:
return None
return {
k: html2text(get_response(v, method='GET').text)
for k, v in urls.items()
}
def rest_api_execute_job(self, job_name: str, e_type: str = "STOP"):
assert e_type in ["START",
"STOP"], "unsupported execute type {}".format(e_type)
return get_response(
"{}/rest-server/api/v1/user/{}/jobs/{}/executionType".format(
self.pai_uri, self.user, job_name),
headers={
'Authorization': 'Bearer {}'.format(self.token),
},
body={
"value": e_type
},
method="PUT",
allowed_status=[200, 202],
).json()
def rest_api_execute_job(self, job_name: str, e_type: str = "STOP"):
assert e_type in ["START",
"STOP"], "unsupported execute type {}".format(e_type)
return get_response(
'PUT',
[
self.rest_srv, 'v1', 'user', self.user, 'jobs', job_name,
'executionType'
],
headers={
'Authorization': 'Bearer {}'.format(self.token),
},
body={
"value": e_type
},
allowed_status=[200, 202],
).json()
def rest_api_job_info(self,
job_name: str = None,
info: str = None,
user: str = None):
import json
import yaml
user = self.user if user is None else user
assert info in [None, 'config',
'ssh'], ('unsupported query information', info)
response = get_response(
'GET', [self.rest_srv, 'v1', 'user', user, 'jobs', job_name, info])
try:
return response.json()
except json.decoder.JSONDecodeError:
return yaml.load(response.text, Loader=yaml.FullLoader)
else:
raise RestSrvError
def connect_jupyter_interactive(self, status: dict = None):
"get the url of notebook if ready"
status = na(status, self.status())
state = self.state(status)
url = None
if state == "RUNNING":
log_url = self.get_logs_url(status)["stderr"]
job_log = html2text(get_response(log_url,
method='GET').text).split('\n')
for line in job_log:
if re.search("The Jupyter Notebook is running at:", line):
container = status["taskRoles"]["main"]["taskStatuses"][0]
url = "http://{}:{}/notebooks/{}".format(
container["containerIp"],
container["containerPorts"]["jupyter"],
self.param("notebook_file") + ".ipynb",
)
break
return dict(state=state, notebook=url)
def rest_api_virtual_clusters(self):
return get_response('{}/rest-server/api/v1/virtual-clusters'.format(
self.pai_uri),
method='GET',
allowed_status=[200]).json()
def rest_api_cluster_info(self):
"refer to https://github.com/microsoft/pai/pull/3281/"
return get_response('GET', [self.rest_srv, 'v1'],
allowed_status=[200]).json()
def rest_api_job_list(self, user: str = None):
return get_response('GET',
[self.rest_srv, 'v1',
('user', user), 'jobs']).json()
