484 lines
20 KiB
Python
484 lines
20 KiB
Python
from datetime import datetime
|
|
from typing import List, Optional, Dict, Any, Tuple
|
|
import random
|
|
import json
|
|
import os
|
|
import shutil
|
|
import asyncio
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.logging import get_logger
|
|
from app.core.config import settings
|
|
from app.db.models.base_entity import LineageNode, LineageEdge
|
|
from app.db.models.ratio_task import RatioInstance, RatioRelation
|
|
from app.db.models import Dataset, DatasetFiles
|
|
from app.db.session import AsyncSessionLocal
|
|
from app.module.dataset.schema.dataset_file import DatasetFileTag
|
|
from app.module.shared.common.lineage import LineageService
|
|
from app.module.shared.schema import TaskStatus, NodeType, EdgeType
|
|
from app.module.ratio.schema.ratio_task import FilterCondition
|
|
from app.module.shared.util.resource_utils import get_concurrent_for_ratio_copy, get_system_resource_info
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class RatioTaskService:
|
|
"""Service for Ratio Task DB operations."""
|
|
|
|
def __init__(self, db: AsyncSession):
|
|
self.db = db
|
|
|
|
async def create_task(
|
|
self,
|
|
*,
|
|
name: str,
|
|
description: Optional[str],
|
|
totals: int,
|
|
config: List[Dict[str, Any]],
|
|
target_dataset_id: Optional[str] = None,
|
|
) -> RatioInstance:
|
|
"""Create a ratio task instance and its relations.
|
|
|
|
config item format: {"dataset_id": str, "counts": int, "filter_conditions": str}
|
|
"""
|
|
logger.info(f"Creating ratio task: name={name}, totals={totals}, items={len(config or [])}")
|
|
|
|
instance = RatioInstance(
|
|
name=name,
|
|
description=description,
|
|
totals=totals,
|
|
target_dataset_id=target_dataset_id,
|
|
status="PENDING",
|
|
)
|
|
self.db.add(instance)
|
|
await self.db.flush() # populate instance.id
|
|
|
|
for item in config or []:
|
|
relation = RatioRelation(
|
|
ratio_instance_id=instance.id,
|
|
source_dataset_id=item.get("dataset_id"),
|
|
counts=int(item.get("counts", 0)),
|
|
filter_conditions=json.dumps({
|
|
'date_range': item.get("filter_conditions").date_range,
|
|
'label': {
|
|
"label":item.get("filter_conditions").label.label,
|
|
"value":item.get("filter_conditions").label.value,
|
|
} if item.get("filter_conditions").label else None,
|
|
})
|
|
)
|
|
logger.info(f"Relation created: {relation.id}, {relation}, {item}, {config}")
|
|
self.db.add(relation)
|
|
|
|
await self.db.commit()
|
|
await self.db.refresh(instance)
|
|
logger.info(f"Ratio task created: {instance.id}")
|
|
return instance
|
|
|
|
# ========================= Execution (Background) ========================= #
|
|
|
|
@staticmethod
|
|
async def execute_dataset_ratio_task(instance_id: str) -> None:
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
inst_res = await session.execute(select(RatioInstance).where(RatioInstance.id == instance_id))
|
|
instance: Optional[RatioInstance] = inst_res.scalar_one_or_none()
|
|
if not instance:
|
|
logger.error(f"Ratio instance not found: {instance_id}")
|
|
return
|
|
|
|
logger.info(f"Starting ratio task {instance_id}")
|
|
logger.info(f"System resources: {get_system_resource_info()}")
|
|
|
|
rel_res = await session.execute(
|
|
select(RatioRelation).where(RatioRelation.ratio_instance_id == instance_id)
|
|
)
|
|
relations: List[RatioRelation] = list(rel_res.scalars().all())
|
|
|
|
instance.status = TaskStatus.RUNNING.name
|
|
|
|
ds_res = await session.execute(select(Dataset).where(Dataset.id == instance.target_dataset_id))
|
|
target_ds: Optional[Dataset] = ds_res.scalar_one_or_none()
|
|
if not target_ds:
|
|
logger.error(f"Target dataset not found for instance {instance_id}")
|
|
instance.status = TaskStatus.FAILED.name
|
|
return
|
|
|
|
max_concurrent = get_concurrent_for_ratio_copy(settings)
|
|
logger.info(f"Using {max_concurrent} concurrent workers for ratio task {instance_id}")
|
|
|
|
added_count, added_size = await RatioTaskService.handle_ratio_relations_parallel(
|
|
relations, session, target_ds, max_concurrent
|
|
)
|
|
|
|
target_ds.file_count = (target_ds.file_count or 0) + added_count
|
|
target_ds.size_bytes = (target_ds.size_bytes or 0) + added_size
|
|
if (target_ds.file_count or 0) > 0:
|
|
target_ds.status = "ACTIVE"
|
|
|
|
instance.status = TaskStatus.COMPLETED.name
|
|
logger.info(f"Ratio task completed: {instance_id}, files={added_count}, size={added_size}")
|
|
|
|
await RatioTaskService._add_task_to_graph(
|
|
session=session,
|
|
src_relations=relations,
|
|
task=instance,
|
|
dst_dataset=target_ds,
|
|
)
|
|
except Exception as e:
|
|
logger.exception(f"Ratio task failed for {instance_id}: {e}")
|
|
try:
|
|
inst_res = await session.execute(select(RatioInstance).where(RatioInstance.id == instance_id))
|
|
instance = inst_res.scalar_one_or_none()
|
|
if instance:
|
|
instance.status = TaskStatus.FAILED.name
|
|
finally:
|
|
pass
|
|
finally:
|
|
await session.commit()
|
|
|
|
@staticmethod
|
|
async def handle_ratio_relations_parallel(
|
|
relations: list[RatioRelation],
|
|
session: AsyncSession,
|
|
target_ds: Dataset,
|
|
max_concurrent: int = 10
|
|
) -> Tuple[int, int]:
|
|
existing_path_rows = await session.execute(
|
|
select(DatasetFiles.file_path).where(DatasetFiles.dataset_id == target_ds.id)
|
|
)
|
|
existing_paths = set(p for p in existing_path_rows.scalars().all() if p)
|
|
source_paths = set()
|
|
|
|
all_copy_tasks: List[Tuple[DatasetFiles, str, str]] = []
|
|
|
|
for rel in relations:
|
|
if not rel.source_dataset_id or not rel.counts or rel.counts <= 0:
|
|
continue
|
|
|
|
files = await RatioTaskService.get_files(rel, session)
|
|
if not files:
|
|
continue
|
|
|
|
pick_n = min(rel.counts or 0, len(files))
|
|
chosen = random.sample(files, pick_n) if pick_n < len(files) else files
|
|
|
|
for file in chosen:
|
|
if file.file_path in source_paths:
|
|
continue
|
|
|
|
dst_prefix = f"/dataset/{target_ds.id}/"
|
|
file_name = RatioTaskService.get_new_file_name(dst_prefix, existing_paths, file)
|
|
new_path = dst_prefix + file_name
|
|
|
|
file_record = DatasetFiles(
|
|
dataset_id=target_ds.id,
|
|
file_name=file_name,
|
|
file_path=new_path,
|
|
file_type=file.file_type,
|
|
file_size=file.file_size,
|
|
check_sum=file.check_sum,
|
|
tags=file.tags,
|
|
tags_updated_at=datetime.now(),
|
|
dataset_filemetadata=file.dataset_filemetadata,
|
|
status="ACTIVE",
|
|
)
|
|
|
|
all_copy_tasks.append((file_record, file.file_path, new_path))
|
|
existing_paths.add(new_path)
|
|
source_paths.add(file.file_path)
|
|
|
|
if not all_copy_tasks:
|
|
return 0, 0
|
|
|
|
dst_dir = f"/dataset/{target_ds.id}/"
|
|
await asyncio.to_thread(os.makedirs, dst_dir, exist_ok=True)
|
|
|
|
semaphore = asyncio.Semaphore(max_concurrent)
|
|
successful_records: List[DatasetFiles] = []
|
|
added_count = 0
|
|
added_size = 0
|
|
|
|
async def copy_with_semaphore(
|
|
file_record: DatasetFiles,
|
|
src_path: str,
|
|
dst_path: str
|
|
) -> Tuple[bool, DatasetFiles]:
|
|
async with semaphore:
|
|
try:
|
|
file_dst_dir = os.path.dirname(dst_path)
|
|
if file_dst_dir != dst_dir:
|
|
await asyncio.to_thread(os.makedirs, file_dst_dir, exist_ok=True)
|
|
|
|
try:
|
|
await asyncio.to_thread(os.link, src_path, dst_path)
|
|
except OSError:
|
|
try:
|
|
await asyncio.to_thread(os.symlink, src_path, dst_path)
|
|
except OSError:
|
|
await asyncio.to_thread(shutil.copy2, src_path, dst_path)
|
|
|
|
return True, file_record
|
|
except Exception as e:
|
|
logger.error(f"Copy failed: {src_path} -> {dst_path}: {e}")
|
|
return False, file_record
|
|
|
|
tasks = [copy_with_semaphore(rec, src, dst) for rec, src, dst in all_copy_tasks]
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
for i, result in enumerate(results):
|
|
if isinstance(result, Exception):
|
|
logger.error(f"Copy task {i} raised exception: {result}")
|
|
continue
|
|
success, file_record = result
|
|
if success:
|
|
added_count += 1
|
|
added_size += int(file_record.file_size or 0)
|
|
successful_records.append(file_record)
|
|
|
|
if successful_records:
|
|
session.add_all(successful_records)
|
|
await session.flush()
|
|
|
|
logger.info(
|
|
f"Parallel copy completed: {added_count}/{len(all_copy_tasks)} files, "
|
|
f"{added_size} bytes, {len(results) - added_count} failures"
|
|
)
|
|
|
|
return added_count, added_size
|
|
|
|
@staticmethod
|
|
async def handle_ratio_relations(relations: list[RatioRelation], session, target_ds: Dataset) -> tuple[int, int]:
|
|
# Preload existing target file paths for deduplication
|
|
existing_path_rows = await session.execute(
|
|
select(DatasetFiles.file_path).where(DatasetFiles.dataset_id == target_ds.id)
|
|
)
|
|
existing_paths = set(p for p in existing_path_rows.scalars().all() if p)
|
|
source_paths = set()
|
|
|
|
added_count = 0
|
|
added_size = 0
|
|
|
|
for rel in relations:
|
|
if not rel.source_dataset_id or not rel.counts or rel.counts <= 0:
|
|
continue
|
|
|
|
files = await RatioTaskService.get_files(rel, session)
|
|
|
|
if not files:
|
|
continue
|
|
|
|
pick_n = min(rel.counts or 0, len(files))
|
|
chosen = random.sample(files, pick_n) if pick_n < len(files) else files
|
|
|
|
# Copy into target dataset with de-dup by target path
|
|
for file in chosen:
|
|
if file.file_path in source_paths:
|
|
continue
|
|
await RatioTaskService.handle_selected_file(existing_paths, file, session, target_ds)
|
|
source_paths.add(file.file_path)
|
|
added_count += 1
|
|
added_size += int(file.file_size or 0)
|
|
|
|
# Periodically flush to avoid huge transactions
|
|
await session.flush()
|
|
return added_count, added_size
|
|
|
|
@staticmethod
|
|
async def handle_selected_file(existing_paths: set[Any], f, session, target_ds: Dataset):
|
|
src_path = f.file_path
|
|
dst_prefix = f"/dataset/{target_ds.id}/"
|
|
file_name = RatioTaskService.get_new_file_name(dst_prefix, existing_paths, f)
|
|
|
|
new_path = dst_prefix + file_name
|
|
dst_dir = os.path.dirname(new_path)
|
|
await asyncio.to_thread(os.makedirs, dst_dir, exist_ok=True)
|
|
await asyncio.to_thread(shutil.copy2, src_path, new_path)
|
|
|
|
file_data = {
|
|
"dataset_id": target_ds.id, # type: ignore
|
|
"file_name": file_name,
|
|
"file_path": new_path,
|
|
"file_type": f.file_type,
|
|
"file_size": f.file_size,
|
|
"check_sum": f.check_sum,
|
|
"tags": f.tags,
|
|
"tags_updated_at": datetime.now(),
|
|
"dataset_filemetadata": f.dataset_filemetadata,
|
|
"status": "ACTIVE",
|
|
}
|
|
file_record = {k: v for k, v in file_data.items() if v is not None}
|
|
session.add(DatasetFiles(**file_record))
|
|
existing_paths.add(new_path)
|
|
|
|
@staticmethod
|
|
def get_new_file_name(dst_prefix: str, existing_paths: set[Any], f) -> str:
|
|
file_name = f.file_name
|
|
new_path = dst_prefix + file_name
|
|
|
|
# Handle file path conflicts by appending a number to the filename
|
|
if new_path in existing_paths:
|
|
file_name_base, file_ext = os.path.splitext(file_name)
|
|
counter = 1
|
|
original_file_name = file_name
|
|
while new_path in existing_paths:
|
|
file_name = f"{file_name_base}_{counter}{file_ext}"
|
|
new_path = f"{dst_prefix}{file_name}"
|
|
counter += 1
|
|
if counter > 1000: # Safety check to prevent infinite loops
|
|
logger.error(f"Could not find unique filename for {original_file_name} after 1000 attempts")
|
|
break
|
|
return file_name
|
|
|
|
@staticmethod
|
|
async def get_files(rel: RatioRelation, session) -> list[Any]:
|
|
# Fetch all files for the source dataset (ACTIVE only)
|
|
files_res = await session.execute(
|
|
select(DatasetFiles).where(
|
|
DatasetFiles.dataset_id == rel.source_dataset_id,
|
|
DatasetFiles.status == "ACTIVE",
|
|
)
|
|
)
|
|
files = list(files_res.scalars().all())
|
|
|
|
# TAG mode: filter by tags according to relation.filter_conditions
|
|
conditions = RatioTaskService._parse_conditions(rel.filter_conditions)
|
|
if conditions:
|
|
files = [f for f in files if RatioTaskService._filter_file(f, conditions)]
|
|
return files
|
|
|
|
# ------------------------- helpers for TAG filtering ------------------------- #
|
|
|
|
@staticmethod
|
|
def _parse_conditions(conditions: Optional[str]) -> Optional[FilterCondition]:
|
|
"""Parse filter_conditions JSON string into a FilterCondition object.
|
|
|
|
Args:
|
|
conditions: JSON string containing filter conditions
|
|
|
|
Returns:
|
|
FilterCondition object if conditions is not None/empty, otherwise None
|
|
"""
|
|
if not conditions:
|
|
return None
|
|
try:
|
|
data = json.loads(conditions)
|
|
return FilterCondition(**data)
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"Failed to parse filter conditions: {e}")
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"Error creating FilterCondition: {e}")
|
|
return None
|
|
|
|
@staticmethod
|
|
def _filter_file(file: DatasetFiles, conditions: FilterCondition) -> bool:
|
|
if not conditions:
|
|
return True
|
|
logger.info(f"start filter file: {file}, conditions: {conditions}")
|
|
|
|
# Check data range condition if provided
|
|
if conditions.date_range and len(conditions.date_range) == 2:
|
|
try:
|
|
from datetime import datetime, timedelta
|
|
start_at = datetime.fromisoformat(conditions.date_range[0])
|
|
end_at = datetime.fromisoformat(conditions.date_range[1])
|
|
if file.tags_updated_at and (file.tags_updated_at < start_at or file.tags_updated_at > end_at):
|
|
return False
|
|
except (ValueError, TypeError) as e:
|
|
logger.warning(f"Invalid data_range value: {conditions.date_range}", e)
|
|
return False
|
|
|
|
# Check label condition if provided
|
|
if conditions.label:
|
|
tags = file.tags
|
|
if not tags:
|
|
return False
|
|
try:
|
|
# tags could be a list of strings or list of objects with 'name'
|
|
all_tags = RatioTaskService.get_all_tags(tags)
|
|
for tag in all_tags:
|
|
if conditions.label.label and tag.get("label") != conditions.label.label:
|
|
continue
|
|
if conditions.label.value is None or len(conditions.label.value) == 0:
|
|
return True
|
|
if tag.get("value") == conditions.label.value:
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
logger.exception(f"Failed to get tags for {file}", e)
|
|
return False
|
|
|
|
return True
|
|
|
|
@staticmethod
|
|
def get_all_tags(tags) -> list[dict]:
|
|
"""获取所有处理后的标签字符串列表"""
|
|
all_tags = list()
|
|
if not tags:
|
|
return all_tags
|
|
|
|
file_tags = []
|
|
for tag_data in tags:
|
|
# 处理可能的命名风格转换（下划线转驼峰）
|
|
processed_data = {}
|
|
for key, value in tag_data.items():
|
|
# 将驼峰转为下划线以匹配 Pydantic 模型字段
|
|
processed_data[key] = value
|
|
# 创建 DatasetFileTag 对象
|
|
file_tag = DatasetFileTag(**processed_data)
|
|
file_tags.append(file_tag)
|
|
|
|
for file_tag in file_tags:
|
|
for tag_data in file_tag.get_tags():
|
|
all_tags.append(tag_data)
|
|
return all_tags
|
|
|
|
@staticmethod
|
|
async def _add_task_to_graph(
|
|
session: AsyncSession,
|
|
src_relations: List[RatioRelation],
|
|
task: RatioInstance,
|
|
dst_dataset: Dataset,
|
|
) -> None:
|
|
"""
|
|
在比例抽取任务完成后，将数据集加入血缘图。
|
|
ratio_task(DATASOURCE) -> dataset(DATASET) via DATA_RATIO edge
|
|
"""
|
|
try:
|
|
if not src_relations:
|
|
logger.warning("Source ratio relations is empty when building lineage graph")
|
|
return
|
|
|
|
lineage_service = LineageService(db=session)
|
|
dst_node = LineageNode(
|
|
id=dst_dataset.id,
|
|
node_type=NodeType.DATASET.value,
|
|
name=dst_dataset.name,
|
|
description=dst_dataset.description,
|
|
)
|
|
for rel in src_relations:
|
|
ds: Optional[Dataset] = await session.get(Dataset, rel.source_dataset_id)
|
|
src_node = LineageNode(
|
|
id=rel.source_dataset_id,
|
|
node_type=NodeType.DATASET.value,
|
|
name=ds.name,
|
|
description=ds.description,
|
|
)
|
|
ratio_edge = LineageEdge(
|
|
process_id=task.id,
|
|
name=task.name,
|
|
edge_type=EdgeType.DATA_RATIO.value,
|
|
description=task.description,
|
|
from_node_id=rel.source_dataset_id,
|
|
to_node_id=dst_node.id,
|
|
)
|
|
await lineage_service.generate_graph(src_node, ratio_edge, dst_node)
|
|
logger.info("Add dataset lineage graph: %s -> %s -> %s", src_node.id, ratio_edge.id, dst_node.id)
|
|
await session.commit()
|
|
logger.info("Add dataset lineage graph success")
|
|
except Exception as exc:
|
|
logger.error("Failed to add dataset lineage graph: %s", exc)
|
|
await session.rollback()
|