Proxy cache service

Dockerfile

@@ -0,0 +1,5 @@
+FROM python:3.11-slim
+WORKDIR /app
+COPY app.py /app/
+EXPOSE 7860
+CMD ["python3", "app.py"]

app.py

@@ -0,0 +1,186 @@
+import http.server
+import json
+import os
+import time
+import shutil
+import threading
+import hashlib
+from pathlib import Path
+from urllib.parse import urlparse, parse_qs
+
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+CACHE_DIR = Path("/tmp/proxy_cache")
+CACHE_DIR.mkdir(exist_ok=True)
+CACHE_DURATION = 5 * 3600  # 5 hours
+MAX_CACHE_SIZE = 2 * 1024 * 1024 * 1024  # 2GB
+
+class SmartCache:
+    def __init__(self):
+        self.lock = threading.Lock()
+        self._start_cleaner()
+
+    def _start_cleaner(self):
+        def clean_loop():
+            while True:
+                time.sleep(300)
+                self.cleanup()
+        t = threading.Thread(target=clean_loop, daemon=True)
+        t.start()
+
+    def get(self, key):
+        fpath = CACHE_DIR / key
+        meta = CACHE_DIR / f"{key}.meta"
+        if fpath.exists() and meta.exists():
+            with open(meta) as f:
+                m = json.load(f)
+            if time.time() - m["cached_at"] < CACHE_DURATION:
+                m["access_count"] = m.get("access_count", 0) + 1
+                with open(meta, "w") as f:
+                    json.dump(m, f)
+                return fpath
+            else:
+                fpath.unlink(missing_ok=True)
+                meta.unlink(missing_ok=True)
+        return None
+
+    def put(self, key, src_path):
+        with self.lock:
+            fpath = CACHE_DIR / key
+            shutil.copy2(src_path, fpath)
+            meta = CACHE_DIR / f"{key}.meta"
+            with open(meta, "w") as f:
+                json.dump({"cached_at": time.time(), "access_count": 0, "size": os.path.getsize(fpath)}, f)
+
+    def cleanup(self):
+        with self.lock:
+            total = sum(f.stat().st_size for f in CACHE_DIR.iterdir() if f.is_file() and not f.name.endswith(".meta"))
+            if total < MAX_CACHE_SIZE:
+                return
+
+            files = []
+            for f in CACHE_DIR.iterdir():
+                if f.is_file() and not f.name.endswith(".meta"):
+                    meta = CACHE_DIR / f"{f.name}.meta"
+                    access = 0
+                    if meta.exists():
+                        with open(meta) as mf:
+                            access = json.load(mf).get("access_count", 0)
+                    files.append((access, f.stat().st_mtime, f))
+            files.sort()
+            removed = 0
+            for _, _, f in files:
+                if total < MAX_CACHE_SIZE * 0.8:
+                    break
+                sz = f.stat().st_size
+                f.unlink()
+                (CACHE_DIR / f"{f.name}.meta").unlink(missing_ok=True)
+                total -= sz
+                removed += 1
+
+cache = SmartCache()
+
+class Handler(http.server.BaseHTTPRequestHandler):
+    def do_GET(self):
+        parsed = urlparse(self.path)
+        params = parse_qs(parsed.query)
+
+        if parsed.path == "/health":
+            total = sum(f.stat().st_size for f in CACHE_DIR.iterdir() if f.is_file() and not f.name.endswith(".meta"))
+            self._json(200, {
+                "status": "ok",
+                "cached_files": len(list(CACHE_DIR.glob("*"))) // 2,
+                "cache_size": total,
+                "cache_max": MAX_CACHE_SIZE,
+            })
+            return
+
+        if parsed.path.startswith("/stream/"):
+            key = parsed.path.split("/stream/")[1]
+            if not key:
+                self._json(400, {"error": "Missing file key"})
+                return
+
+            cached = cache.get(key)
+            if cached:
+                self._serve_file(cached, key)
+            else:
+                self._json(404, {"error": "File not in cache. Use /cache to preload."})
+            return
+
+        self._json(404, {"error": "Not found"})
+
+    def do_POST(self):
+        length = int(self.headers.get("Content-Length", 0))
+        body = json.loads(self.rfile.read(length)) if length else {}
+
+        if self.path == "/cache":
+            file_key = body.get("key", "")
+            file_url = body.get("url", "")
+            if not file_key or not file_url:
+                self._json(400, {"error": "Missing key or url"})
+                return
+            try:
+                import urllib.request
+                temp = CACHE_DIR / f"dl_{int(time.time())}"
+                urllib.request.urlretrieve(file_url, temp)
+                cache.put(file_key, temp)
+                temp.unlink()
+                self._json(200, {"status": "cached", "key": file_key})
+            except Exception as e:
+                self._json(500, {"error": str(e)})
+            return
+
+        if self.path == "/preload":
+            dataset = body.get("dataset", "")
+            file_name = body.get("file_name", "")
+            if not dataset or not file_name:
+                self._json(400, {"error": "Missing dataset or file_name"})
+                return
+            dl_url = f"https://huggingface.co/datasets/{dataset}/resolve/main/{file_name}"
+            try:
+                import urllib.request
+                req = urllib.request.Request(dl_url, headers={"Authorization": f"Bearer {HF_TOKEN}"})
+                temp = CACHE_DIR / f"dl_{int(time.time())}"
+                urllib.request.urlretrieve(dl_url, temp)
+                key = f"{dataset}/{file_name}"
+                cache.put(key, temp)
+                temp.unlink()
+                self._json(200, {"status": "preloaded", "key": key})
+            except Exception as e:
+                self._json(500, {"error": str(e)})
+            return
+
+        if self.path == "/cleanup":
+            for f in CACHE_DIR.iterdir():
+                if f.is_file():
+                    f.unlink()
+            self._json(200, {"status": "cleaned"})
+            return
+
+        self._json(404, {"error": "Not found"})
+
+    def _serve_file(self, fpath, filename):
+        self.send_response(200)
+        self.send_header("Content-Type", "video/mp4")
+        self.send_header("Content-Disposition", f'inline; filename="{filename}"')
+        self.send_header("Accept-Ranges", "bytes")
+        self.send_header("Cache-Control", f"public, max-age={CACHE_DURATION}")
+        self.end_headers()
+        with open(fpath, "rb") as f:
+            shutil.copyfileobj(f, self.wfile)
+
+    def _json(self, code, data):
+        self.send_response(code)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Access-Control-Allow-Origin", "*")
+        self.end_headers()
+        self.wfile.write(json.dumps(data).encode())
+
+    def log_message(self, format, *args):
+        pass
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    httpd = http.server.HTTPServer(("0.0.0.0", port), Handler)
+    print(f"Proxy running on port {port}")
+    httpd.serve_forever()