def publish_task( cls, task_id: str, company_id: str, publish_model: bool, force: bool, status_reason: str = "", status_message: str = "", ) -> dict: task = cls.get_task_with_access( task_id, company_id=company_id, requires_write_access=True ) if not force: validate_status_change(task.status, TaskStatus.published) previous_task_status = task.status output = task.output or Output() publish_failed = False try: # set state to publishing task.status = TaskStatus.publishing task.save() # publish task models if task.output.model and publish_model: output_model = ( Model.objects(id=task.output.model) .only("id", "task", "ready") .first() ) if output_model and not output_model.ready: cls.model_set_ready( model_id=task.output.model, company_id=company_id, publish_task=False, ) # set task status to published, and update (or set) it's new output (view and models) return ChangeStatusRequest( task=task, new_status=TaskStatus.published, force=force, status_reason=status_reason, status_message=status_message, ).execute(published=datetime.utcnow(), output=output) except Exception as ex: publish_failed = True raise ex finally: if publish_failed: task.status = previous_task_status task.save()
def update(call): assert isinstance(call, APICall) identity = call.identity model_id = call.data["model"] force = call.data.get("force", False) with translate_errors_context(): query = dict(id=model_id, company=identity.company) model = Model.objects(**query).only("id", "task").first() if not model: raise errors.bad_request.InvalidModelId(**query) deleted_model_id = f"__DELETED__{model_id}" using_tasks = Task.objects(execution__model=model_id).only("id") if using_tasks: if not force: raise errors.bad_request.ModelInUse( "as execution model, use force=True to delete", num_tasks=len(using_tasks), ) # update deleted model id in using tasks using_tasks.update( execution__model=deleted_model_id, upsert=False, multi=True ) if model.task: task = Task.objects(id=model.task).first() if task and task.status == TaskStatus.published: if not force: raise errors.bad_request.ModelCreatingTaskExists( "and published, use force=True to delete", task=model.task ) task.update( output__model=deleted_model_id, output__error=f"model deleted on {datetime.utcnow().isoformat()}", upsert=False, ) del_count = Model.objects(**query).delete() call.result.data = dict(deleted=del_count > 0)
def validate_execution_model(task, allow_only_public=False): if not task.execution or not task.execution.model: return company = None if allow_only_public else task.company model_id = task.execution.model model = Model.objects( Q(id=model_id) & get_company_or_none_constraint(company)).first() if not model: raise errors.bad_request.InvalidModelId(model=model_id) return model
def get_all(call: APICall, company_id, _): conform_tag_fields(call, call.data) with translate_errors_context(): with TimingContext("mongo", "models_get_all"): models = Model.get_many( company=company_id, parameters=call.data, query_dict=call.data, allow_public=True, ) conform_output_tags(call, models) call.result.data = {"models": models}
def get_all_ex(call: APICall): conform_tag_fields(call, call.data) with translate_errors_context(): with TimingContext("mongo", "models_get_all_ex"): models = Model.get_many_with_join( company=call.identity.company, query_dict=call.data, allow_public=True, query_options=get_all_query_options, ) conform_output_tags(call, models) call.result.data = {"models": models}
def get_all_ex(call): assert isinstance(call, APICall) with translate_errors_context(): with TimingContext("mongo", "models_get_all_ex"): models = Model.get_many_with_join( company=call.identity.company, query_dict=call.data, allow_public=True, query_options=get_all_query_options, ) call.result.data = {"models": models}
def edit(call: APICall, company_id, _): model_id = call.data["model"] with translate_errors_context(): query = dict(id=model_id, company=company_id) model = Model.objects(**query).first() if not model: raise errors.bad_request.InvalidModelId(**query) fields = parse_model_fields(call, create_fields) fields = prepare_update_fields(call, company_id, fields) for key in fields: field = getattr(model, key, None) value = fields[key] if (field and isinstance(value, dict) and isinstance(field, EmbeddedDocument)): d = field.to_mongo(use_db_field=False).to_dict() d.update(value) fields[key] = d iteration = call.data.get("iteration") task_id = model.task or fields.get("task") if task_id and iteration is not None: TaskBLL.update_statistics( task_id=task_id, company_id=company_id, last_iteration_max=iteration, ) if fields: updated = model.update(upsert=False, **fields) if updated: new_project = fields.get("project", model.project) if new_project != model.project: _reset_cached_tags(company_id, projects=[new_project, model.project]) else: _update_cached_tags(company_id, project=model.project, fields=fields) conform_output_tags(call, fields) call.result.data_model = UpdateResponse(updated=updated, fields=fields) else: call.result.data_model = UpdateResponse(updated=0)
def get_by_id(call: APICall, company_id, _): model_id = call.data["model"] with translate_errors_context(): models = Model.get_many( company=company_id, query_dict=call.data, query=Q(id=model_id), allow_public=True, ) if not models: raise errors.bad_request.InvalidModelId( "no such public or company model", id=model_id, company=company_id, ) conform_output_tags(call, models[0]) call.result.data = {"model": models[0]}
def _resolve_entities( cls, experiments: Sequence[str] = None, projects: Sequence[str] = None, task_statuses: Sequence[str] = None, ) -> Dict[Type[mongoengine.Document], Set[mongoengine.Document]]: entities = defaultdict(set) if projects: print("Reading projects...") entities[Project].update(cls._resolve_type(Project, projects)) print("--> Reading project experiments...") query = Q( project__in=list( set(filter(None, (p.id for p in entities[Project])))), system_tags__nin=[EntityVisibility.archived.value], ) if task_statuses: query &= Q(status__in=list(set(task_statuses))) objs = Task.objects(query) entities[Task].update(o for o in objs if o.id not in (experiments or [])) if experiments: print("Reading experiments...") entities[Task].update(cls._resolve_type(Task, experiments)) print("--> Reading experiments projects...") objs = Project.objects(id__in=list( set(filter(None, (p.project for p in entities[Task]))))) project_ids = {p.id for p in entities[Project]} entities[Project].update(o for o in objs if o.id not in project_ids) model_ids = { model_id for task in entities[Task] for model_id in (task.output.model, task.execution.model) if model_id } if model_ids: print("Reading models...") entities[Model] = set(Model.objects(id__in=list(model_ids))) return entities
def get_outputs_for_deletion(task, force=False): with TimingContext("mongo", "get_task_models"): models = TaskOutputs( attrgetter("ready"), Model, Model.objects(task=task.id).only("id", "task", "ready"), ) if not force and models.published: raise errors.bad_request.TaskCannotBeDeleted( "has output models, use force=True", task=task.id, models=len(models.published), ) if task.output.model: output_model = get_output_model(task, force) if output_model: if output_model.ready: models.published.append(output_model) else: models.draft.append(output_model) if models.draft: with TimingContext("mongo", "get_execution_models"): model_ids = [m.id for m in models.draft] dependent_tasks = Task.objects( execution__model__in=model_ids).only("id", "execution.model") busy_models = [t.execution.model for t in dependent_tasks] models.draft[:] = [ m for m in models.draft if m.id not in busy_models ] with TimingContext("mongo", "get_task_children"): tasks = Task.objects(parent=task.id).only("id", "parent", "status") published_tasks = [ task for task in tasks if task.status == TaskStatus.published ] if not force and published_tasks: raise errors.bad_request.TaskCannotBeDeleted( "has children, use force=True", task=task.id, children=published_tasks) return models, tasks
def get_by_id(call): assert isinstance(call, APICall) model_id = call.data["model"] with translate_errors_context(): res = Model.get_many( company=call.identity.company, query_dict=call.data, query=Q(id=model_id), allow_public=True, ) if not res: raise errors.bad_request.InvalidModelId( "no such public or company model", id=model_id, company=call.identity.company, ) call.result.data = {"model": res[0]}
def edit(call): assert isinstance(call, APICall) identity = call.identity model_id = call.data["model"] with translate_errors_context(): query = dict(id=model_id, company=identity.company) model = Model.objects(**query).first() if not model: raise errors.bad_request.InvalidModelId(**query) fields = parse_model_fields(call, create_fields) fields = prepare_update_fields(call, fields) for key in fields: field = getattr(model, key, None) value = fields[key] if (field and isinstance(value, dict) and isinstance(field, EmbeddedDocument)): d = field.to_mongo(use_db_field=False).to_dict() d.update(value) fields[key] = d iteration = call.data.get("iteration") task_id = model.task or fields.get('task') if task_id and iteration is not None: TaskBLL.update_statistics( task_id=task_id, company_id=identity.company, last_iteration_max=iteration, ) if fields: updated = model.update(upsert=False, **fields) call.result.data_model = UpdateResponse(updated=updated, fields=fields) else: call.result.data_model = UpdateResponse(updated=0)
def make_public(call: APICall, company_id, request: MakePublicRequest): with translate_errors_context(): call.result.data = Model.set_public(company_id, request.ids, invalid_cls=InvalidModelId, enabled=False)
def update_for_task(call: APICall, company_id, _): task_id = call.data["task"] uri = call.data.get("uri") iteration = call.data.get("iteration") override_model_id = call.data.get("override_model_id") if not (uri or override_model_id) or (uri and override_model_id): raise errors.bad_request.MissingRequiredFields( "exactly one field is required", fields=("uri", "override_model_id")) with translate_errors_context(): query = dict(id=task_id, company=company_id) task = Task.get_for_writing( id=task_id, company=company_id, _only=["output", "execution", "name", "status", "project"], ) if not task: raise errors.bad_request.InvalidTaskId(**query) allowed_states = [TaskStatus.created, TaskStatus.in_progress] if task.status not in allowed_states: raise errors.bad_request.InvalidTaskStatus( f"model can only be updated for tasks in the {allowed_states} states", **query, ) if override_model_id: query = dict(company=company_id, id=override_model_id) model = Model.objects(**query).first() if not model: raise errors.bad_request.InvalidModelId(**query) else: if "name" not in call.data: # use task name if name not provided call.data["name"] = task.name if "comment" not in call.data: call.data[ "comment"] = f"Created by task `{task.name}` ({task.id})" if task.output and task.output.model: # model exists, update res = _update_model(call, company_id, model_id=task.output.model).to_struct() res.update({"id": task.output.model, "created": False}) call.result.data = res return # new model, create fields = parse_model_fields(call, create_fields) # create and save model model = Model( id=database.utils.id(), created=datetime.utcnow(), user=call.identity.user, company=company_id, project=task.project, framework=task.execution.framework, parent=task.execution.model, design=task.execution.model_desc, labels=task.execution.model_labels, ready=(task.status == TaskStatus.published), **fields, ) model.save() _update_cached_tags(company_id, project=model.project, fields=fields) TaskBLL.update_statistics( task_id=task_id, company_id=company_id, last_iteration_max=iteration, output__model=model.id, ) call.result.data = {"id": model.id, "created": True}
def parse_model_fields(call, valid_fields): fields = parse_from_call(call.data, valid_fields, Model.get_fields()) conform_tag_fields(call, fields, validate=True) return fields
def parse_model_fields(call, valid_fields): fields = parse_from_call(call.data, valid_fields, Model.get_fields()) tags = fields.get("tags") if tags: fields["tags"] = list(set(tags)) return fields
parse_from_call, get_company_or_none_constraint, filter_fields, ) from service_repo import APICall, endpoint from services.utils import conform_tag_fields, conform_output_tags from timing_context import TimingContext log = config.logger(__file__) get_all_query_options = Model.QueryParameterOptions( pattern_fields=("name", "comment"), fields=("ready",), list_fields=( "tags", "system_tags", "framework", "uri", "id", "project", "task", "parent", ), ) @endpoint("models.get_by_id", required_fields=["model"]) def get_by_id(call): assert isinstance(call, APICall) model_id = call.data["model"] with translate_errors_context(): models = Model.get_many(
def _cleanup_model(cls, model: Model): model.company = "" model.user = "" model.tags = cls._filter_out_export_tags(model.tags)