mirror of
https://github.com/clearml/clearml-server
synced 2025-01-31 19:06:55 +00:00
276 lines
8.7 KiB
Python
276 lines
8.7 KiB
Python
import csv
|
|
from collections import defaultdict
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from io import StringIO
|
|
from operator import itemgetter
|
|
from typing import Mapping, Type, Sequence, Optional, Callable
|
|
|
|
from flask import stream_with_context
|
|
from mongoengine import Q
|
|
|
|
from apiserver.apierrors import errors
|
|
from apiserver.apimodels.organization import (
|
|
TagsRequest,
|
|
EntitiesCountRequest,
|
|
DownloadForGetAll,
|
|
EntityType,
|
|
PrepareDownloadForGetAll,
|
|
)
|
|
from apiserver.bll.model import Metadata
|
|
from apiserver.bll.organization import OrgBLL, Tags
|
|
from apiserver.bll.project import ProjectBLL
|
|
from apiserver.config_repo import config
|
|
from apiserver.database.model import User, AttributedDocument, EntityVisibility
|
|
from apiserver.database.model.model import Model
|
|
from apiserver.database.model.project import Project
|
|
from apiserver.database.model.task.task import Task, TaskType
|
|
from apiserver.redis_manager import redman
|
|
from apiserver.service_repo import endpoint, APICall
|
|
from apiserver.services.models import conform_model_data
|
|
from apiserver.services.tasks import (
|
|
escape_execution_parameters,
|
|
_hidden_query,
|
|
conform_task_data,
|
|
)
|
|
from apiserver.services.utils import get_tags_filter_dictionary, sort_tags_response
|
|
from apiserver.utilities import json
|
|
from apiserver.utilities.dicts import nested_get
|
|
|
|
org_bll = OrgBLL()
|
|
project_bll = ProjectBLL()
|
|
redis = redman.connection("apiserver")
|
|
|
|
|
|
@endpoint("organization.get_tags", request_data_model=TagsRequest)
|
|
def get_tags(call: APICall, company, request: TagsRequest):
|
|
filter_dict = get_tags_filter_dictionary(request.filter)
|
|
ret = defaultdict(set)
|
|
for entity in Tags.Model, Tags.Task:
|
|
tags = org_bll.get_tags(
|
|
company, entity, include_system=request.include_system, filter_=filter_dict,
|
|
)
|
|
for field, vals in tags.items():
|
|
ret[field] |= vals
|
|
|
|
call.result.data = sort_tags_response(ret)
|
|
|
|
|
|
@endpoint("organization.get_user_companies")
|
|
def get_user_companies(call: APICall, company_id: str, _):
|
|
users = [
|
|
{"id": u.id, "name": u.name, "avatar": u.avatar}
|
|
for u in User.objects(company=company_id).only("avatar", "name", "company")
|
|
]
|
|
|
|
call.result.data = {
|
|
"companies": [
|
|
{
|
|
"id": company_id,
|
|
"name": call.identity.company_name,
|
|
"allocated": len(users),
|
|
"owners": sorted(users, key=itemgetter("name")),
|
|
}
|
|
]
|
|
}
|
|
|
|
|
|
@endpoint("organization.get_entities_count")
|
|
def get_entities_count(call: APICall, company, request: EntitiesCountRequest):
|
|
entity_classes: Mapping[str, Type[AttributedDocument]] = {
|
|
"projects": Project,
|
|
"tasks": Task,
|
|
"models": Model,
|
|
"pipelines": Project,
|
|
"datasets": Project,
|
|
"reports": Task,
|
|
}
|
|
ret = {}
|
|
for field, entity_cls in entity_classes.items():
|
|
data = call.data.get(field)
|
|
if data is None:
|
|
continue
|
|
|
|
if field == "reports":
|
|
data["type"] = TaskType.report
|
|
data["include_subprojects"] = True
|
|
|
|
if request.active_users:
|
|
if entity_cls is Project:
|
|
requested_ids = data.get("id")
|
|
if isinstance(requested_ids, str):
|
|
requested_ids = [requested_ids]
|
|
ids, _ = project_bll.get_projects_with_selected_children(
|
|
company=company,
|
|
users=request.active_users,
|
|
project_ids=requested_ids,
|
|
allow_public=request.allow_public,
|
|
)
|
|
if not ids:
|
|
ret[field] = 0
|
|
continue
|
|
data["id"] = ids
|
|
elif not data.get("user"):
|
|
data["user"] = request.active_users
|
|
|
|
query = Q()
|
|
if (
|
|
entity_cls in (Project, Task)
|
|
and field not in ("reports", "pipelines", "datasets")
|
|
and not request.search_hidden
|
|
):
|
|
query &= Q(system_tags__ne=EntityVisibility.hidden.value)
|
|
|
|
ret[field] = entity_cls.get_count(
|
|
company=company,
|
|
query_dict=data,
|
|
query=query,
|
|
allow_public=request.allow_public,
|
|
)
|
|
|
|
call.result.data = ret
|
|
|
|
|
|
def _get_download_getter_fn(
|
|
company: str,
|
|
call: APICall,
|
|
call_data: dict,
|
|
allow_public: bool,
|
|
entity_type: EntityType,
|
|
) -> Optional[Callable[[int, int], Sequence[dict]]]:
|
|
def get_task_data() -> Sequence[dict]:
|
|
tasks = Task.get_many_with_join(
|
|
company=company,
|
|
query_dict=call_data,
|
|
query=_hidden_query(call_data),
|
|
allow_public=allow_public,
|
|
)
|
|
conform_task_data(call, tasks)
|
|
return tasks
|
|
|
|
def get_model_data() -> Sequence[dict]:
|
|
models = Model.get_many_with_join(
|
|
company=company, query_dict=call_data, allow_public=allow_public,
|
|
)
|
|
conform_model_data(call, models)
|
|
return models
|
|
|
|
if entity_type == EntityType.task:
|
|
call_data = escape_execution_parameters(call_data)
|
|
get_fn = get_task_data
|
|
elif entity_type == EntityType.model:
|
|
call_data = Metadata.escape_query_parameters(call_data)
|
|
get_fn = get_model_data
|
|
else:
|
|
raise errors.bad_request.ValidationError(
|
|
f"Unsupported entity type: {str(entity_type)}"
|
|
)
|
|
|
|
def getter(page: int, page_size: int) -> Sequence[dict]:
|
|
call_data.pop("scroll_id", None)
|
|
call_data.pop("start", None)
|
|
call_data.pop("size", None)
|
|
call_data.pop("refresh_scroll", None)
|
|
call_data["page"] = page
|
|
call_data["page_size"] = page_size
|
|
return get_fn()
|
|
|
|
return getter
|
|
|
|
|
|
@endpoint("organization.prepare_download_for_get_all")
|
|
def prepare_download_for_get_all(
|
|
call: APICall, company: str, request: PrepareDownloadForGetAll
|
|
):
|
|
# validate input params
|
|
getter = _get_download_getter_fn(
|
|
company,
|
|
call,
|
|
call_data=call.data.copy(),
|
|
allow_public=request.allow_public,
|
|
entity_type=request.entity_type,
|
|
)
|
|
if getter:
|
|
getter(0, 1)
|
|
|
|
redis.setex(
|
|
f"get_all_download_{call.id}",
|
|
int(config.get("services.organization.download.redis_timeout_sec", 300)),
|
|
json.dumps(call.data),
|
|
)
|
|
|
|
call.result.data = dict(prepare_id=call.id)
|
|
|
|
|
|
@endpoint("organization.download_for_get_all")
|
|
def download_for_get_all(call: APICall, company, request: DownloadForGetAll):
|
|
request_data = redis.get(f"get_all_download_{request.prepare_id}")
|
|
if not request_data:
|
|
raise errors.bad_request.InvalidId(
|
|
f"prepare ID not found", prepare_id=request.prepare_id
|
|
)
|
|
|
|
try:
|
|
call_data = json.loads(request_data)
|
|
request = PrepareDownloadForGetAll(**call_data)
|
|
except Exception as ex:
|
|
raise errors.server_error.DataError("failed parsing prepared data", ex=ex)
|
|
|
|
class SingleLine:
|
|
def write(self, line: str) -> str:
|
|
return line
|
|
|
|
def generate():
|
|
projection = call_data.get("only_fields", [])
|
|
|
|
get_fn = _get_download_getter_fn(
|
|
company,
|
|
call,
|
|
call_data=call_data,
|
|
allow_public=request.allow_public,
|
|
entity_type=request.entity_type,
|
|
)
|
|
if not get_fn:
|
|
yield csv.writer(SingleLine()).writerow(projection)
|
|
return
|
|
|
|
fields = [path.split(".") for path in projection]
|
|
|
|
def get_entity_field_as_str(data: dict, field: Sequence[str]) -> str:
|
|
val = nested_get(data, field, "")
|
|
if isinstance(val, dict):
|
|
val = val.get("id", "")
|
|
|
|
return str(val)
|
|
|
|
def get_projected_fields(data: dict) -> Sequence[str]:
|
|
return [get_entity_field_as_str(data, f) for f in fields]
|
|
|
|
with ThreadPoolExecutor(1) as pool:
|
|
page = 0
|
|
page_size = int(
|
|
config.get("services.organization.download.batch_size", 500)
|
|
)
|
|
future = pool.submit(get_fn, page, page_size)
|
|
|
|
while True:
|
|
result = future.result()
|
|
if not result:
|
|
break
|
|
|
|
page += 1
|
|
future = pool.submit(get_fn, page, page_size)
|
|
|
|
with StringIO() as fp:
|
|
writer = csv.writer(fp)
|
|
if page == 1:
|
|
writer.writerow(projection)
|
|
writer.writerows(get_projected_fields(r) for r in result)
|
|
yield fp.getvalue()
|
|
|
|
if page == 0:
|
|
yield csv.writer(SingleLine()).writerow(projection)
|
|
|
|
call.result.filename = f"{request.entity_type}_export.csv"
|
|
call.result.content_type = "text/csv"
|
|
call.result.raw_data = stream_with_context(generate())
|