def download_dataset(name, path, dataset_type, url, auth_headers):
"""download dataset
Download dataset object from S3 through ZumoLabs backend.
Args:
name (str): name of dataset to download
path (str): output directory
dataset_type (str): type of dataset to download
url (str): backend endpoint
auth_headers: authentication for backend
Returns:
str: output file path
"""
endpoint = f"{url}/api/v1/{dataset_type}-data-sets/"
r = requests.get(endpoint, params={"name": name}, headers=auth_headers)
if r.status_code != 200:
r.raise_for_status()
response = json.loads(r.text)
if response["count"] != 1:
raise NameError(f"found {response['count']} datasets for name {name}")
dataset = response["results"][0]
endpoint = (
f"{url}/api/v1/{dataset['dataset_type']}-data-sets/{dataset['id']}/download"
)
r = requests.get(endpoint, headers=auth_headers)
if r.status_code != 200:
r.raise_for_status()
response = json.loads(r.text)
name_slug = f"{dataset['name'].replace(' ', '_')}-{dataset['id'][:8]}.zip"
output_path = to_pathlib_path(path) / name_slug
download_url(response["redirect_link"], output_path)
return output_path
def fetch_logs(resource, name, path, url, auth_headers):
"""fetch logs
Fetch LOG_TYPES for a backend run.
Args:
type (str): resource to fetch logs
name (str): name of resource
path (str): output_dir
url (str): backend endpoint
auth_headers: authentication for backend
"""
endpoint = f"{url}/api/v1/{resource}/"
r = requests.get(endpoint, params={"name": name}, headers=auth_headers)
if r.status_code != 200:
r.raise_for_status()
response = json.loads(r.text)
if response["count"] != 1:
raise NameError(
f"found {response['count']} {resource} for name {name}")
obj = response["results"][0]
endpoint = f"{url}/api/v1/{resource}/{obj['id']}/logs"
r = requests.get(endpoint, headers=auth_headers)
if r.status_code != 200:
r.raise_for_status()
response = json.loads(r.text)
output_dir = to_pathlib_path(path)
for log_type in LOG_TYPES:
output_path = output_dir / f"{log_type}.log"
download_url(response[log_type]["redirect_link"], output_path)
def fetch_dataset(name, path, dataset_type, url, token):
""" fetch a dataset from ragnarok """
endpoint = f'{url}/api/v1/{dataset_type}-data-sets/'
params = {'name': name}
r = requests.get(endpoint, params=params, headers=auth_headers(token))
if r.status_code != 200:
log.warning(f'Unable to fetch {dataset_type} datasets')
return
response = json.loads(r.text)
if response['count'] != 1:
log.warning(f'Unable to find {dataset_type} dataset with name "{name}"')
return
dataset = response['results'][0]
endpoint = f"{url}/api/v1/{dataset['dataset_type']}-data-sets/{dataset['id']}/download"
r = requests.get(endpoint, headers=auth_headers(token))
if r.status_code != 200:
log.warning(f"Unable to get download link for dataset {dataset['id']}")
return
response = json.loads(r.text)
name_slug = f"{dataset['name'].replace(' ', '_')}-{dataset['id'][:8]}.zip"
output_path = to_pathlib_path(path) / name_slug
download_url(response['redirect_link'], output_path)
def fetch_scene(name, path, url, token):
""" fetch a scene from ragnarok """
endpoint = f'{url}/api/v1/scenes/'
params = {'name': name}
r = requests.get(endpoint, params=params, headers=auth_headers(token))
if r.status_code != 200:
log.warning(f'Unable to fetch scenes')
return
response = json.loads(r.text)
if response['count'] != 1:
log.warning(f'Unable to find scene with name "{name}"')
return
scene = response['results'][0]
endpoint = f"{url}/api/v1/scenes/{scene['id']}/download"
r = requests.get(endpoint, headers=auth_headers(token))
if r.status_code != 200:
log.warning(f"Unable to get download link for scene {scene['id']}")
return
response = json.loads(r.text)
name_slug = f"{scene['name'].replace(' ', '_')}-{scene['id'][:8]}.zip"
output_path = to_pathlib_path(path) / name_slug
download_url(response['redirect_link'], output_path)