Spaces:
Paused
Paused
File size: 3,623 Bytes
81d5a2e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import os
import shutil
from pathlib import Path
from typing import Dict, Any, List
def resolve_storage_dir() -> Path:
# Check if /data is mounted and writable (HF Spaces persistent storage bucket)
candidate = Path("/data")
if candidate.exists():
try:
test_file = candidate / ".write_test"
test_file.touch(exist_ok=True)
test_file.unlink(missing_ok=True)
return candidate
except Exception:
pass
# Local fallback
local_data = Path("./data").resolve()
local_data.mkdir(parents=True, exist_ok=True)
return local_data
STORAGE_DIR = resolve_storage_dir()
# Ensure essential directories
RUNS_DIR = STORAGE_DIR / "runs"
ARTIFACTS_DIR = STORAGE_DIR / "artifacts"
MEMORY_DIR = STORAGE_DIR / "memory"
for d in (RUNS_DIR, ARTIFACTS_DIR, MEMORY_DIR):
d.mkdir(parents=True, exist_ok=True)
def get_storage_info() -> Dict[str, Any]:
try:
total, used, free = shutil.disk_usage(STORAGE_DIR)
is_hf_bucket = str(STORAGE_DIR).startswith("/data")
file_count = sum(1 for _ in STORAGE_DIR.rglob("*") if _.is_file())
return {
"path": str(STORAGE_DIR),
"is_hf_bucket": is_hf_bucket,
"total_gb": round(total / (1024**3), 2),
"used_gb": round(used / (1024**3), 2),
"free_gb": round(free / (1024**3), 2),
"percent_used": round((used / total) * 100, 1) if total > 0 else 0,
"file_count": file_count,
}
except Exception as e:
return {
"path": str(STORAGE_DIR),
"is_hf_bucket": False,
"error": str(e),
"total_gb": 0,
"used_gb": 0,
"free_gb": 0,
"percent_used": 0,
"file_count": 0,
}
def list_storage_files(subpath: str = "") -> List[Dict[str, Any]]:
target = (STORAGE_DIR / subpath.lstrip("/")).resolve()
if not str(target).startswith(str(STORAGE_DIR)):
raise ValueError("Path traversal denied.")
if not target.exists():
return []
results = []
for item in sorted(target.iterdir(), key=lambda p: (not p.is_dir(), p.name)):
rel = item.relative_to(STORAGE_DIR).as_posix()
results.append({
"name": item.name,
"path": rel,
"is_dir": item.is_dir(),
"size_bytes": item.stat().st_size if item.is_file() else 0,
"modified_time": item.stat().st_mtime,
})
return results
def read_storage_file(rel_path: str) -> str:
target = (STORAGE_DIR / rel_path.lstrip("/")).resolve()
if not str(target).startswith(str(STORAGE_DIR)):
raise ValueError("Path traversal denied.")
if not target.is_file():
raise FileNotFoundError(f"File not found: {rel_path}")
return target.read_text(encoding="utf-8", errors="replace")
def write_storage_file(rel_path: str, content: str) -> str:
target = (STORAGE_DIR / rel_path.lstrip("/")).resolve()
if not str(target).startswith(str(STORAGE_DIR)):
raise ValueError("Path traversal denied.")
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(content, encoding="utf-8")
return str(target.relative_to(STORAGE_DIR))
def delete_storage_file(rel_path: str) -> bool:
target = (STORAGE_DIR / rel_path.lstrip("/")).resolve()
if not str(target).startswith(str(STORAGE_DIR)):
raise ValueError("Path traversal denied.")
if not target.exists():
return False
if target.is_dir():
shutil.rmtree(target)
else:
target.unlink()
return True
|