560 lines
22 KiB
Python
560 lines
22 KiB
Python
"""
|
|
Artifact Repository
|
|
Repository for model artifact operations
|
|
"""
|
|
|
|
from typing import Optional, List, Dict, Any
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy import select, and_, text, desc
|
|
from datetime import datetime, timedelta
|
|
import structlog
|
|
|
|
from .base import TrainingBaseRepository
|
|
from app.models.training import ModelArtifact
|
|
from shared.database.exceptions import DatabaseError, ValidationError
|
|
|
|
logger = structlog.get_logger()
|
|
|
|
|
|
class ArtifactRepository(TrainingBaseRepository):
|
|
"""Repository for model artifact operations"""
|
|
|
|
def __init__(self, session: AsyncSession, cache_ttl: Optional[int] = 1800):
|
|
# Artifacts are stable, longer cache time (30 minutes)
|
|
super().__init__(ModelArtifact, session, cache_ttl)
|
|
|
|
async def create_artifact(self, artifact_data: Dict[str, Any]) -> ModelArtifact:
|
|
"""Create a new model artifact record"""
|
|
try:
|
|
# Validate artifact data
|
|
validation_result = self._validate_training_data(
|
|
artifact_data,
|
|
["model_id", "tenant_id", "artifact_type", "file_path"]
|
|
)
|
|
|
|
if not validation_result["is_valid"]:
|
|
raise ValidationError(f"Invalid artifact data: {validation_result['errors']}")
|
|
|
|
# Set default values
|
|
if "storage_location" not in artifact_data:
|
|
artifact_data["storage_location"] = "local"
|
|
|
|
# Create artifact record
|
|
artifact = await self.create(artifact_data)
|
|
|
|
logger.info("Model artifact created",
|
|
model_id=artifact.model_id,
|
|
tenant_id=artifact.tenant_id,
|
|
artifact_type=artifact.artifact_type,
|
|
file_path=artifact.file_path)
|
|
|
|
return artifact
|
|
|
|
except ValidationError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error("Failed to create model artifact",
|
|
model_id=artifact_data.get("model_id"),
|
|
error=str(e))
|
|
raise DatabaseError(f"Failed to create artifact: {str(e)}")
|
|
|
|
async def get_artifacts_by_model(
|
|
self,
|
|
model_id: str,
|
|
artifact_type: str = None
|
|
) -> List[ModelArtifact]:
|
|
"""Get all artifacts for a model"""
|
|
try:
|
|
filters = {"model_id": model_id}
|
|
if artifact_type:
|
|
filters["artifact_type"] = artifact_type
|
|
|
|
return await self.get_multi(
|
|
filters=filters,
|
|
order_by="created_at",
|
|
order_desc=True
|
|
)
|
|
except Exception as e:
|
|
logger.error("Failed to get artifacts by model",
|
|
model_id=model_id,
|
|
artifact_type=artifact_type,
|
|
error=str(e))
|
|
raise DatabaseError(f"Failed to get artifacts: {str(e)}")
|
|
|
|
async def get_artifacts_by_tenant(
|
|
self,
|
|
tenant_id: str,
|
|
artifact_type: str = None,
|
|
skip: int = 0,
|
|
limit: int = 100
|
|
) -> List[ModelArtifact]:
|
|
"""Get artifacts for a tenant"""
|
|
try:
|
|
filters = {"tenant_id": tenant_id}
|
|
if artifact_type:
|
|
filters["artifact_type"] = artifact_type
|
|
|
|
return await self.get_multi(
|
|
filters=filters,
|
|
skip=skip,
|
|
limit=limit,
|
|
order_by="created_at",
|
|
order_desc=True
|
|
)
|
|
except Exception as e:
|
|
logger.error("Failed to get artifacts by tenant",
|
|
tenant_id=tenant_id,
|
|
error=str(e))
|
|
raise DatabaseError(f"Failed to get tenant artifacts: {str(e)}")
|
|
|
|
async def get_artifact_by_path(self, file_path: str) -> Optional[ModelArtifact]:
|
|
"""Get artifact by file path"""
|
|
try:
|
|
return await self.get_by_field("file_path", file_path)
|
|
except Exception as e:
|
|
logger.error("Failed to get artifact by path",
|
|
file_path=file_path,
|
|
error=str(e))
|
|
raise DatabaseError(f"Failed to get artifact: {str(e)}")
|
|
|
|
async def update_artifact_size(self, artifact_id: int, file_size_bytes: int) -> Optional[ModelArtifact]:
|
|
"""Update artifact file size"""
|
|
try:
|
|
return await self.update(artifact_id, {"file_size_bytes": file_size_bytes})
|
|
except Exception as e:
|
|
logger.error("Failed to update artifact size",
|
|
artifact_id=artifact_id,
|
|
error=str(e))
|
|
return None
|
|
|
|
async def update_artifact_checksum(self, artifact_id: int, checksum: str) -> Optional[ModelArtifact]:
|
|
"""Update artifact checksum for integrity verification"""
|
|
try:
|
|
return await self.update(artifact_id, {"checksum": checksum})
|
|
except Exception as e:
|
|
logger.error("Failed to update artifact checksum",
|
|
artifact_id=artifact_id,
|
|
error=str(e))
|
|
return None
|
|
|
|
async def mark_artifact_expired(self, artifact_id: int, expires_at: datetime = None) -> Optional[ModelArtifact]:
|
|
"""Mark artifact for expiration/cleanup"""
|
|
try:
|
|
if not expires_at:
|
|
expires_at = datetime.now()
|
|
|
|
return await self.update(artifact_id, {"expires_at": expires_at})
|
|
except Exception as e:
|
|
logger.error("Failed to mark artifact as expired",
|
|
artifact_id=artifact_id,
|
|
error=str(e))
|
|
return None
|
|
|
|
async def get_expired_artifacts(self, days_expired: int = 0) -> List[ModelArtifact]:
|
|
"""Get artifacts that have expired"""
|
|
try:
|
|
cutoff_date = datetime.now() - timedelta(days=days_expired)
|
|
|
|
query_text = """
|
|
SELECT * FROM model_artifacts
|
|
WHERE expires_at IS NOT NULL
|
|
AND expires_at <= :cutoff_date
|
|
ORDER BY expires_at ASC
|
|
"""
|
|
|
|
result = await self.session.execute(text(query_text), {"cutoff_date": cutoff_date})
|
|
|
|
expired_artifacts = []
|
|
for row in result.fetchall():
|
|
record_dict = dict(row._mapping)
|
|
artifact = self.model(**record_dict)
|
|
expired_artifacts.append(artifact)
|
|
|
|
return expired_artifacts
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to get expired artifacts",
|
|
days_expired=days_expired,
|
|
error=str(e))
|
|
return []
|
|
|
|
async def cleanup_expired_artifacts(self, days_expired: int = 0) -> int:
|
|
"""Clean up expired artifacts"""
|
|
try:
|
|
cutoff_date = datetime.now() - timedelta(days=days_expired)
|
|
|
|
query_text = """
|
|
DELETE FROM model_artifacts
|
|
WHERE expires_at IS NOT NULL
|
|
AND expires_at <= :cutoff_date
|
|
"""
|
|
|
|
result = await self.session.execute(text(query_text), {"cutoff_date": cutoff_date})
|
|
deleted_count = result.rowcount
|
|
|
|
logger.info("Cleaned up expired artifacts",
|
|
deleted_count=deleted_count,
|
|
days_expired=days_expired)
|
|
|
|
return deleted_count
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to cleanup expired artifacts",
|
|
days_expired=days_expired,
|
|
error=str(e))
|
|
raise DatabaseError(f"Artifact cleanup failed: {str(e)}")
|
|
|
|
async def get_large_artifacts(self, min_size_mb: int = 100) -> List[ModelArtifact]:
|
|
"""Get artifacts larger than specified size"""
|
|
try:
|
|
min_size_bytes = min_size_mb * 1024 * 1024 # Convert MB to bytes
|
|
|
|
query_text = """
|
|
SELECT * FROM model_artifacts
|
|
WHERE file_size_bytes >= :min_size_bytes
|
|
ORDER BY file_size_bytes DESC
|
|
"""
|
|
|
|
result = await self.session.execute(text(query_text), {"min_size_bytes": min_size_bytes})
|
|
|
|
large_artifacts = []
|
|
for row in result.fetchall():
|
|
record_dict = dict(row._mapping)
|
|
artifact = self.model(**record_dict)
|
|
large_artifacts.append(artifact)
|
|
|
|
return large_artifacts
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to get large artifacts",
|
|
min_size_mb=min_size_mb,
|
|
error=str(e))
|
|
return []
|
|
|
|
async def get_artifacts_by_storage_location(
|
|
self,
|
|
storage_location: str,
|
|
tenant_id: str = None
|
|
) -> List[ModelArtifact]:
|
|
"""Get artifacts by storage location"""
|
|
try:
|
|
filters = {"storage_location": storage_location}
|
|
if tenant_id:
|
|
filters["tenant_id"] = tenant_id
|
|
|
|
return await self.get_multi(
|
|
filters=filters,
|
|
order_by="created_at",
|
|
order_desc=True
|
|
)
|
|
except Exception as e:
|
|
logger.error("Failed to get artifacts by storage location",
|
|
storage_location=storage_location,
|
|
tenant_id=tenant_id,
|
|
error=str(e))
|
|
raise DatabaseError(f"Failed to get artifacts: {str(e)}")
|
|
|
|
async def get_artifact_statistics(self, tenant_id: str = None) -> Dict[str, Any]:
|
|
"""Get artifact statistics"""
|
|
try:
|
|
base_filters = {}
|
|
if tenant_id:
|
|
base_filters["tenant_id"] = tenant_id
|
|
|
|
# Get basic counts
|
|
total_artifacts = await self.count(filters=base_filters)
|
|
|
|
# Get artifacts by type
|
|
type_query_params = {}
|
|
type_query_filter = ""
|
|
if tenant_id:
|
|
type_query_filter = "WHERE tenant_id = :tenant_id"
|
|
type_query_params["tenant_id"] = tenant_id
|
|
|
|
type_query = text(f"""
|
|
SELECT artifact_type, COUNT(*) as count
|
|
FROM model_artifacts
|
|
{type_query_filter}
|
|
GROUP BY artifact_type
|
|
ORDER BY count DESC
|
|
""")
|
|
|
|
result = await self.session.execute(type_query, type_query_params)
|
|
artifacts_by_type = {row.artifact_type: row.count for row in result.fetchall()}
|
|
|
|
# Get storage location stats
|
|
location_query = text(f"""
|
|
SELECT
|
|
storage_location,
|
|
COUNT(*) as count,
|
|
SUM(COALESCE(file_size_bytes, 0)) as total_size_bytes
|
|
FROM model_artifacts
|
|
{type_query_filter}
|
|
GROUP BY storage_location
|
|
ORDER BY count DESC
|
|
""")
|
|
|
|
location_result = await self.session.execute(location_query, type_query_params)
|
|
storage_stats = {}
|
|
total_size_bytes = 0
|
|
|
|
for row in location_result.fetchall():
|
|
storage_stats[row.storage_location] = {
|
|
"artifact_count": row.count,
|
|
"total_size_bytes": int(row.total_size_bytes or 0),
|
|
"total_size_mb": round((row.total_size_bytes or 0) / (1024 * 1024), 2)
|
|
}
|
|
total_size_bytes += row.total_size_bytes or 0
|
|
|
|
# Get expired artifacts count
|
|
expired_artifacts = len(await self.get_expired_artifacts())
|
|
|
|
return {
|
|
"total_artifacts": total_artifacts,
|
|
"expired_artifacts": expired_artifacts,
|
|
"active_artifacts": total_artifacts - expired_artifacts,
|
|
"artifacts_by_type": artifacts_by_type,
|
|
"storage_statistics": storage_stats,
|
|
"total_storage": {
|
|
"total_size_bytes": total_size_bytes,
|
|
"total_size_mb": round(total_size_bytes / (1024 * 1024), 2),
|
|
"total_size_gb": round(total_size_bytes / (1024 * 1024 * 1024), 2)
|
|
}
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to get artifact statistics",
|
|
tenant_id=tenant_id,
|
|
error=str(e))
|
|
return {
|
|
"total_artifacts": 0,
|
|
"expired_artifacts": 0,
|
|
"active_artifacts": 0,
|
|
"artifacts_by_type": {},
|
|
"storage_statistics": {},
|
|
"total_storage": {
|
|
"total_size_bytes": 0,
|
|
"total_size_mb": 0.0,
|
|
"total_size_gb": 0.0
|
|
}
|
|
}
|
|
|
|
async def verify_artifact_integrity(self, artifact_id: int) -> Dict[str, Any]:
|
|
"""Verify artifact file integrity with actual file system checks"""
|
|
try:
|
|
import os
|
|
import hashlib
|
|
|
|
artifact = await self.get_by_id(artifact_id)
|
|
if not artifact:
|
|
return {"exists": False, "error": "Artifact not found"}
|
|
|
|
# Check if file exists
|
|
file_exists = os.path.exists(artifact.file_path)
|
|
if not file_exists:
|
|
return {
|
|
"artifact_id": artifact_id,
|
|
"file_path": artifact.file_path,
|
|
"exists": False,
|
|
"checksum_valid": False,
|
|
"size_valid": False,
|
|
"storage_location": artifact.storage_location,
|
|
"last_verified": datetime.now().isoformat(),
|
|
"error": "File does not exist on disk"
|
|
}
|
|
|
|
# Verify file size
|
|
actual_size = os.path.getsize(artifact.file_path)
|
|
size_valid = True
|
|
if artifact.file_size_bytes:
|
|
size_valid = (actual_size == artifact.file_size_bytes)
|
|
|
|
# Verify checksum if stored
|
|
checksum_valid = True
|
|
actual_checksum = None
|
|
if artifact.checksum:
|
|
# Calculate checksum of actual file
|
|
sha256_hash = hashlib.sha256()
|
|
try:
|
|
with open(artifact.file_path, "rb") as f:
|
|
# Read file in chunks to handle large files
|
|
for byte_block in iter(lambda: f.read(4096), b""):
|
|
sha256_hash.update(byte_block)
|
|
actual_checksum = sha256_hash.hexdigest()
|
|
checksum_valid = (actual_checksum == artifact.checksum)
|
|
except Exception as checksum_error:
|
|
logger.error(f"Failed to calculate checksum: {checksum_error}")
|
|
checksum_valid = False
|
|
actual_checksum = None
|
|
|
|
# Overall integrity status
|
|
integrity_valid = file_exists and size_valid and checksum_valid
|
|
|
|
result = {
|
|
"artifact_id": artifact_id,
|
|
"file_path": artifact.file_path,
|
|
"exists": file_exists,
|
|
"checksum_valid": checksum_valid,
|
|
"size_valid": size_valid,
|
|
"integrity_valid": integrity_valid,
|
|
"storage_location": artifact.storage_location,
|
|
"last_verified": datetime.now().isoformat(),
|
|
"details": {
|
|
"stored_size_bytes": artifact.file_size_bytes,
|
|
"actual_size_bytes": actual_size if file_exists else None,
|
|
"stored_checksum": artifact.checksum,
|
|
"actual_checksum": actual_checksum
|
|
}
|
|
}
|
|
|
|
if not integrity_valid:
|
|
issues = []
|
|
if not file_exists:
|
|
issues.append("file_missing")
|
|
if not size_valid:
|
|
issues.append("size_mismatch")
|
|
if not checksum_valid:
|
|
issues.append("checksum_mismatch")
|
|
result["issues"] = issues
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to verify artifact integrity",
|
|
artifact_id=artifact_id,
|
|
error=str(e))
|
|
return {
|
|
"exists": False,
|
|
"error": f"Verification failed: {str(e)}"
|
|
}
|
|
|
|
async def migrate_artifacts_to_storage(
|
|
self,
|
|
from_location: str,
|
|
to_location: str,
|
|
tenant_id: str = None,
|
|
copy_only: bool = False,
|
|
verify: bool = True
|
|
) -> Dict[str, Any]:
|
|
"""Migrate artifacts from one storage location to another with actual file operations"""
|
|
try:
|
|
import os
|
|
import shutil
|
|
import hashlib
|
|
|
|
# Get artifacts to migrate
|
|
artifacts = await self.get_artifacts_by_storage_location(from_location, tenant_id)
|
|
|
|
migrated_count = 0
|
|
failed_count = 0
|
|
failed_artifacts = []
|
|
verified_count = 0
|
|
|
|
for artifact in artifacts:
|
|
try:
|
|
# Determine new file path
|
|
new_file_path = artifact.file_path.replace(from_location, to_location, 1)
|
|
|
|
# Create destination directory if it doesn't exist
|
|
dest_dir = os.path.dirname(new_file_path)
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
|
|
# Check if source file exists
|
|
if not os.path.exists(artifact.file_path):
|
|
logger.warning(f"Source file not found: {artifact.file_path}")
|
|
failed_count += 1
|
|
failed_artifacts.append({
|
|
"artifact_id": artifact.id,
|
|
"file_path": artifact.file_path,
|
|
"reason": "source_file_not_found"
|
|
})
|
|
continue
|
|
|
|
# Copy or move file
|
|
if copy_only:
|
|
shutil.copy2(artifact.file_path, new_file_path)
|
|
logger.debug(f"Copied file from {artifact.file_path} to {new_file_path}")
|
|
else:
|
|
shutil.move(artifact.file_path, new_file_path)
|
|
logger.debug(f"Moved file from {artifact.file_path} to {new_file_path}")
|
|
|
|
# Verify file was copied/moved successfully
|
|
if verify and os.path.exists(new_file_path):
|
|
# Verify file size
|
|
new_size = os.path.getsize(new_file_path)
|
|
if artifact.file_size_bytes and new_size != artifact.file_size_bytes:
|
|
logger.warning(f"File size mismatch after migration: {new_file_path}")
|
|
failed_count += 1
|
|
failed_artifacts.append({
|
|
"artifact_id": artifact.id,
|
|
"file_path": new_file_path,
|
|
"reason": "size_mismatch_after_migration"
|
|
})
|
|
continue
|
|
|
|
# Verify checksum if available
|
|
if artifact.checksum:
|
|
sha256_hash = hashlib.sha256()
|
|
with open(new_file_path, "rb") as f:
|
|
for byte_block in iter(lambda: f.read(4096), b""):
|
|
sha256_hash.update(byte_block)
|
|
new_checksum = sha256_hash.hexdigest()
|
|
|
|
if new_checksum != artifact.checksum:
|
|
logger.warning(f"Checksum mismatch after migration: {new_file_path}")
|
|
failed_count += 1
|
|
failed_artifacts.append({
|
|
"artifact_id": artifact.id,
|
|
"file_path": new_file_path,
|
|
"reason": "checksum_mismatch_after_migration"
|
|
})
|
|
continue
|
|
|
|
verified_count += 1
|
|
|
|
# Update database with new location
|
|
await self.update(artifact.id, {
|
|
"storage_location": to_location,
|
|
"file_path": new_file_path
|
|
})
|
|
|
|
migrated_count += 1
|
|
|
|
except Exception as migration_error:
|
|
logger.error("Failed to migrate artifact",
|
|
artifact_id=artifact.id,
|
|
error=str(migration_error))
|
|
failed_count += 1
|
|
failed_artifacts.append({
|
|
"artifact_id": artifact.id,
|
|
"file_path": artifact.file_path,
|
|
"reason": str(migration_error)
|
|
})
|
|
|
|
logger.info("Artifact migration completed",
|
|
from_location=from_location,
|
|
to_location=to_location,
|
|
migrated_count=migrated_count,
|
|
failed_count=failed_count,
|
|
verified_count=verified_count)
|
|
|
|
return {
|
|
"from_location": from_location,
|
|
"to_location": to_location,
|
|
"total_artifacts": len(artifacts),
|
|
"migrated_count": migrated_count,
|
|
"failed_count": failed_count,
|
|
"verified_count": verified_count if verify else None,
|
|
"success_rate": round((migrated_count / len(artifacts)) * 100, 2) if artifacts else 100,
|
|
"copy_only": copy_only,
|
|
"failed_artifacts": failed_artifacts if failed_artifacts else None
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to migrate artifacts",
|
|
from_location=from_location,
|
|
to_location=to_location,
|
|
error=str(e))
|
|
return {
|
|
"error": f"Migration failed: {str(e)}"
|
|
} |