Commit
·
56a64e0
1
Parent(s):
151773c
0.0.2.8 V Beta
Browse files- LoadBalancer.py +51 -63
- app.py +20 -215
- main.py +219 -0
- requirements.txt +3 -2
LoadBalancer.py
CHANGED
@@ -1,11 +1,10 @@
|
|
1 |
import os
|
2 |
import json
|
3 |
-
|
|
|
4 |
import re
|
|
|
5 |
from tvdb import fetch_and_cache_json
|
6 |
-
from threading import Event, Thread
|
7 |
-
import time
|
8 |
-
import logging
|
9 |
from utils import convert_to_gb
|
10 |
from api import InstancesAPI
|
11 |
|
@@ -15,13 +14,13 @@ download_progress = {}
|
|
15 |
|
16 |
class LoadBalancer:
|
17 |
def __init__(self, cache_dir, index_file, token, repo, polling_interval=4, max_retries=3, initial_delay=1):
|
18 |
-
self.version = "0.0.2.
|
19 |
self.instances = []
|
20 |
self.instances_health = {}
|
21 |
self.polling_interval = polling_interval
|
22 |
self.max_retries = max_retries
|
23 |
self.initial_delay = initial_delay
|
24 |
-
self.stop_event = Event()
|
25 |
self.instances_api = InstancesAPI(self.instances)
|
26 |
self.CACHE_DIR = cache_dir
|
27 |
self.INDEX_FILE = index_file
|
@@ -42,16 +41,11 @@ class LoadBalancer:
|
|
42 |
# Load the file structure JSON
|
43 |
self.load_file_structure()
|
44 |
|
45 |
-
# Start polling and file checking in separate
|
46 |
-
|
47 |
-
|
48 |
-
polling_thread.start()
|
49 |
-
|
50 |
-
file_checking_thread = Thread(target=self.check_file_updates)
|
51 |
-
file_checking_thread.daemon = True
|
52 |
-
file_checking_thread.start()
|
53 |
|
54 |
-
def load_file_structure(self):
|
55 |
if not os.path.exists(self.INDEX_FILE):
|
56 |
raise FileNotFoundError(f"{self.INDEX_FILE} not found. Please make sure the file exists.")
|
57 |
|
@@ -59,23 +53,21 @@ class LoadBalancer:
|
|
59 |
self.file_structure = json.load(f)
|
60 |
logging.info("File structure loaded successfully.")
|
61 |
|
62 |
-
def check_file_updates(self):
|
63 |
while not self.stop_event.is_set():
|
64 |
if self.index_file_last_modified != os.path.getmtime(self.INDEX_FILE):
|
65 |
logging.info(f"{self.INDEX_FILE} has been updated. Re-indexing...")
|
66 |
indexer() # Re-run the indexer
|
67 |
-
self.load_file_structure() # Reload the file structure
|
68 |
self.index_file_last_modified = os.path.getmtime(self.INDEX_FILE)
|
69 |
|
70 |
-
# Restart prefetching
|
71 |
-
if hasattr(self, '
|
72 |
-
self.
|
73 |
|
74 |
-
self.
|
75 |
-
self.prefetch_thread.daemon = True
|
76 |
-
self.prefetch_thread.start()
|
77 |
|
78 |
-
|
79 |
|
80 |
def register_instance(self, instance_url):
|
81 |
if instance_url not in self.instances:
|
@@ -92,8 +84,8 @@ class LoadBalancer:
|
|
92 |
else:
|
93 |
logging.info(f"Instance {instance_url} not found for removal.")
|
94 |
|
95 |
-
def get_reports(self):
|
96 |
-
reports = self.instances_api.fetch_reports()
|
97 |
|
98 |
# Initialize temporary JSON data holders
|
99 |
temp_film_store = {}
|
@@ -137,30 +129,27 @@ class LoadBalancer:
|
|
137 |
logging.info("Film and TV Stores processed successfully.")
|
138 |
self.update_instances_health(instance=instance_url, cache_size=cache_size)
|
139 |
|
140 |
-
def start_polling(self):
|
141 |
logging.info("Starting polling.")
|
142 |
while not self.stop_event.is_set():
|
143 |
-
self.get_reports()
|
144 |
-
|
145 |
logging.info("Polling stopped.")
|
146 |
|
147 |
-
def stop_polling(self):
|
148 |
logging.info("Stopping polling.")
|
149 |
self.stop_event.set()
|
150 |
|
151 |
-
def start_prefetching(self):
|
152 |
-
"""Start the metadata prefetching
|
153 |
-
self.prefetch_metadata()
|
154 |
-
|
155 |
-
#################################################################
|
156 |
|
157 |
def update_instances_health(self, instance, cache_size):
|
158 |
-
self.instances_health[instance] = {"used":cache_size["cache_size"],
|
159 |
"total": "50 GB"}
|
160 |
logging.info(f"Updated instance {instance} with cache size {cache_size}")
|
161 |
|
162 |
-
|
163 |
-
def download_film_to_best_instance(self, title):
|
164 |
"""
|
165 |
Downloads a film to the first instance that has more free space on the self.instance_health list variable.
|
166 |
The instance_health looks like this:
|
@@ -187,14 +176,14 @@ class LoadBalancer:
|
|
187 |
best_instance = instance_url
|
188 |
|
189 |
if best_instance:
|
190 |
-
result = self.instances_api.download_film(best_instance, title)
|
191 |
film_id = result["film_id"]
|
192 |
status = result["status"]
|
193 |
progress_url = f'{best_instance}/api/progress/{film_id}'
|
194 |
response = {
|
195 |
-
"film_id":film_id,
|
196 |
-
"status":status,
|
197 |
-
"progress_url":progress_url
|
198 |
}
|
199 |
|
200 |
return response
|
@@ -202,9 +191,9 @@ class LoadBalancer:
|
|
202 |
logging.error("No suitable instance found for downloading the film.")
|
203 |
return {"error": "No suitable instance found for downloading the film."}
|
204 |
|
205 |
-
def download_episode_to_best_instance(self, title, season, episode):
|
206 |
"""
|
207 |
-
Downloads
|
208 |
The instance_health looks like this:
|
209 |
{
|
210 |
"https://unicone-studio-instance1.hf.space": {
|
@@ -213,9 +202,9 @@ class LoadBalancer:
|
|
213 |
}
|
214 |
}
|
215 |
Args:
|
216 |
-
title (str): The title of the
|
217 |
-
season (str): The season of the
|
218 |
-
episode (str): The
|
219 |
"""
|
220 |
best_instance = None
|
221 |
max_free_space = -1
|
@@ -231,23 +220,22 @@ class LoadBalancer:
|
|
231 |
best_instance = instance_url
|
232 |
|
233 |
if best_instance:
|
234 |
-
result = self.instances_api.download_episode(best_instance, title, season, episode)
|
235 |
episode_id = result["episode_id"]
|
236 |
status = result["status"]
|
237 |
progress_url = f'{best_instance}/api/progress/{episode_id}'
|
238 |
response = {
|
239 |
-
"episode_id":episode_id,
|
240 |
-
"status":status,
|
241 |
-
"progress_url":progress_url
|
242 |
}
|
243 |
|
244 |
return response
|
245 |
else:
|
246 |
-
logging.error("No suitable instance found for downloading the
|
247 |
-
return {"error": "No suitable instance found for downloading the
|
248 |
|
249 |
-
|
250 |
-
def find_movie_path(self, title):
|
251 |
"""Find the path of the movie in the JSON data based on the title."""
|
252 |
for directory in self.file_structure:
|
253 |
if directory['type'] == 'directory' and directory['path'] == 'films':
|
@@ -258,7 +246,7 @@ class LoadBalancer:
|
|
258 |
return item['path']
|
259 |
return None
|
260 |
|
261 |
-
def find_tv_path(self, title):
|
262 |
"""Find the path of the TV show in the JSON data based on the title."""
|
263 |
for directory in self.file_structure:
|
264 |
if directory['type'] == 'directory' and directory['path'] == 'tv':
|
@@ -267,7 +255,7 @@ class LoadBalancer:
|
|
267 |
return sub_directory['path']
|
268 |
return None
|
269 |
|
270 |
-
def get_tv_structure(self, title):
|
271 |
"""Find the path of the TV show in the JSON data based on the title."""
|
272 |
for directory in self.file_structure:
|
273 |
if directory['type'] == 'directory' and directory['path'] == 'tv':
|
@@ -276,11 +264,11 @@ class LoadBalancer:
|
|
276 |
return sub_directory
|
277 |
return None
|
278 |
|
279 |
-
def get_film_id(self, title):
|
280 |
"""Generate a film ID based on the title."""
|
281 |
return title.replace(" ", "_").lower()
|
282 |
|
283 |
-
def prefetch_metadata(self):
|
284 |
"""Prefetch metadata for all items in the file structure."""
|
285 |
for item in self.file_structure:
|
286 |
if 'contents' in item:
|
@@ -303,9 +291,9 @@ class LoadBalancer:
|
|
303 |
title = parts[0].strip()
|
304 |
year = int(parts[-1])
|
305 |
|
306 |
-
fetch_and_cache_json(original_title, title, media_type, year)
|
307 |
|
308 |
-
def get_all_tv_shows(self):
|
309 |
"""Get all TV shows from the indexed cache structure JSON file."""
|
310 |
tv_shows = {}
|
311 |
for directory in self.file_structure:
|
@@ -326,7 +314,7 @@ class LoadBalancer:
|
|
326 |
})
|
327 |
return tv_shows
|
328 |
|
329 |
-
def get_all_films(self):
|
330 |
"""Get all films from the indexed cache structure JSON file."""
|
331 |
films = []
|
332 |
for directory in self.file_structure:
|
@@ -334,4 +322,4 @@ class LoadBalancer:
|
|
334 |
for sub_directory in directory['contents']:
|
335 |
if sub_directory['type'] == 'directory':
|
336 |
films.append(sub_directory['path'])
|
337 |
-
return films
|
|
|
1 |
import os
|
2 |
import json
|
3 |
+
import asyncio
|
4 |
+
import logging
|
5 |
import re
|
6 |
+
from indexer import indexer
|
7 |
from tvdb import fetch_and_cache_json
|
|
|
|
|
|
|
8 |
from utils import convert_to_gb
|
9 |
from api import InstancesAPI
|
10 |
|
|
|
14 |
|
15 |
class LoadBalancer:
|
16 |
def __init__(self, cache_dir, index_file, token, repo, polling_interval=4, max_retries=3, initial_delay=1):
|
17 |
+
self.version = "0.0.2.8 V Beta"
|
18 |
self.instances = []
|
19 |
self.instances_health = {}
|
20 |
self.polling_interval = polling_interval
|
21 |
self.max_retries = max_retries
|
22 |
self.initial_delay = initial_delay
|
23 |
+
self.stop_event = asyncio.Event()
|
24 |
self.instances_api = InstancesAPI(self.instances)
|
25 |
self.CACHE_DIR = cache_dir
|
26 |
self.INDEX_FILE = index_file
|
|
|
41 |
# Load the file structure JSON
|
42 |
self.load_file_structure()
|
43 |
|
44 |
+
# Start polling and file checking in separate tasks
|
45 |
+
asyncio.create_task(self.start_polling())
|
46 |
+
asyncio.create_task(self.check_file_updates())
|
|
|
|
|
|
|
|
|
|
|
47 |
|
48 |
+
async def load_file_structure(self):
|
49 |
if not os.path.exists(self.INDEX_FILE):
|
50 |
raise FileNotFoundError(f"{self.INDEX_FILE} not found. Please make sure the file exists.")
|
51 |
|
|
|
53 |
self.file_structure = json.load(f)
|
54 |
logging.info("File structure loaded successfully.")
|
55 |
|
56 |
+
async def check_file_updates(self):
|
57 |
while not self.stop_event.is_set():
|
58 |
if self.index_file_last_modified != os.path.getmtime(self.INDEX_FILE):
|
59 |
logging.info(f"{self.INDEX_FILE} has been updated. Re-indexing...")
|
60 |
indexer() # Re-run the indexer
|
61 |
+
await self.load_file_structure() # Reload the file structure
|
62 |
self.index_file_last_modified = os.path.getmtime(self.INDEX_FILE)
|
63 |
|
64 |
+
# Restart prefetching task
|
65 |
+
if hasattr(self, 'prefetch_task') and not self.prefetch_task.done():
|
66 |
+
await self.prefetch_task
|
67 |
|
68 |
+
self.prefetch_task = asyncio.create_task(self.start_prefetching())
|
|
|
|
|
69 |
|
70 |
+
await asyncio.sleep(120) # Check every 2 minutes
|
71 |
|
72 |
def register_instance(self, instance_url):
|
73 |
if instance_url not in self.instances:
|
|
|
84 |
else:
|
85 |
logging.info(f"Instance {instance_url} not found for removal.")
|
86 |
|
87 |
+
async def get_reports(self):
|
88 |
+
reports = await self.instances_api.fetch_reports()
|
89 |
|
90 |
# Initialize temporary JSON data holders
|
91 |
temp_film_store = {}
|
|
|
129 |
logging.info("Film and TV Stores processed successfully.")
|
130 |
self.update_instances_health(instance=instance_url, cache_size=cache_size)
|
131 |
|
132 |
+
async def start_polling(self):
|
133 |
logging.info("Starting polling.")
|
134 |
while not self.stop_event.is_set():
|
135 |
+
await self.get_reports()
|
136 |
+
await asyncio.sleep(self.polling_interval)
|
137 |
logging.info("Polling stopped.")
|
138 |
|
139 |
+
async def stop_polling(self):
|
140 |
logging.info("Stopping polling.")
|
141 |
self.stop_event.set()
|
142 |
|
143 |
+
async def start_prefetching(self):
|
144 |
+
"""Start the metadata prefetching."""
|
145 |
+
await self.prefetch_metadata()
|
|
|
|
|
146 |
|
147 |
def update_instances_health(self, instance, cache_size):
|
148 |
+
self.instances_health[instance] = {"used": cache_size["cache_size"],
|
149 |
"total": "50 GB"}
|
150 |
logging.info(f"Updated instance {instance} with cache size {cache_size}")
|
151 |
|
152 |
+
async def download_film_to_best_instance(self, title):
|
|
|
153 |
"""
|
154 |
Downloads a film to the first instance that has more free space on the self.instance_health list variable.
|
155 |
The instance_health looks like this:
|
|
|
176 |
best_instance = instance_url
|
177 |
|
178 |
if best_instance:
|
179 |
+
result = await self.instances_api.download_film(best_instance, title)
|
180 |
film_id = result["film_id"]
|
181 |
status = result["status"]
|
182 |
progress_url = f'{best_instance}/api/progress/{film_id}'
|
183 |
response = {
|
184 |
+
"film_id": film_id,
|
185 |
+
"status": status,
|
186 |
+
"progress_url": progress_url
|
187 |
}
|
188 |
|
189 |
return response
|
|
|
191 |
logging.error("No suitable instance found for downloading the film.")
|
192 |
return {"error": "No suitable instance found for downloading the film."}
|
193 |
|
194 |
+
async def download_episode_to_best_instance(self, title, season, episode):
|
195 |
"""
|
196 |
+
Downloads an episode to the first instance that has more free space on the self.instance_health list variable.
|
197 |
The instance_health looks like this:
|
198 |
{
|
199 |
"https://unicone-studio-instance1.hf.space": {
|
|
|
202 |
}
|
203 |
}
|
204 |
Args:
|
205 |
+
title (str): The title of the TV show.
|
206 |
+
season (str): The season of the TV show.
|
207 |
+
episode (str): The episode of the TV show.
|
208 |
"""
|
209 |
best_instance = None
|
210 |
max_free_space = -1
|
|
|
220 |
best_instance = instance_url
|
221 |
|
222 |
if best_instance:
|
223 |
+
result = await self.instances_api.download_episode(best_instance, title, season, episode)
|
224 |
episode_id = result["episode_id"]
|
225 |
status = result["status"]
|
226 |
progress_url = f'{best_instance}/api/progress/{episode_id}'
|
227 |
response = {
|
228 |
+
"episode_id": episode_id,
|
229 |
+
"status": status,
|
230 |
+
"progress_url": progress_url
|
231 |
}
|
232 |
|
233 |
return response
|
234 |
else:
|
235 |
+
logging.error("No suitable instance found for downloading the episode.")
|
236 |
+
return {"error": "No suitable instance found for downloading the episode."}
|
237 |
|
238 |
+
async def find_movie_path(self, title):
|
|
|
239 |
"""Find the path of the movie in the JSON data based on the title."""
|
240 |
for directory in self.file_structure:
|
241 |
if directory['type'] == 'directory' and directory['path'] == 'films':
|
|
|
246 |
return item['path']
|
247 |
return None
|
248 |
|
249 |
+
async def find_tv_path(self, title):
|
250 |
"""Find the path of the TV show in the JSON data based on the title."""
|
251 |
for directory in self.file_structure:
|
252 |
if directory['type'] == 'directory' and directory['path'] == 'tv':
|
|
|
255 |
return sub_directory['path']
|
256 |
return None
|
257 |
|
258 |
+
async def get_tv_structure(self, title):
|
259 |
"""Find the path of the TV show in the JSON data based on the title."""
|
260 |
for directory in self.file_structure:
|
261 |
if directory['type'] == 'directory' and directory['path'] == 'tv':
|
|
|
264 |
return sub_directory
|
265 |
return None
|
266 |
|
267 |
+
async def get_film_id(self, title):
|
268 |
"""Generate a film ID based on the title."""
|
269 |
return title.replace(" ", "_").lower()
|
270 |
|
271 |
+
async def prefetch_metadata(self):
|
272 |
"""Prefetch metadata for all items in the file structure."""
|
273 |
for item in self.file_structure:
|
274 |
if 'contents' in item:
|
|
|
291 |
title = parts[0].strip()
|
292 |
year = int(parts[-1])
|
293 |
|
294 |
+
await fetch_and_cache_json(original_title, title, media_type, year)
|
295 |
|
296 |
+
async def get_all_tv_shows(self):
|
297 |
"""Get all TV shows from the indexed cache structure JSON file."""
|
298 |
tv_shows = {}
|
299 |
for directory in self.file_structure:
|
|
|
314 |
})
|
315 |
return tv_shows
|
316 |
|
317 |
+
async def get_all_films(self):
|
318 |
"""Get all films from the indexed cache structure JSON file."""
|
319 |
films = []
|
320 |
for directory in self.file_structure:
|
|
|
322 |
for sub_directory in directory['contents']:
|
323 |
if sub_directory['type'] == 'directory':
|
324 |
films.append(sub_directory['path'])
|
325 |
+
return films
|
app.py
CHANGED
@@ -1,217 +1,22 @@
|
|
1 |
-
|
2 |
-
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
10 |
|
11 |
-
|
12 |
-
app = Flask(__name__)
|
13 |
-
CORS(app)
|
14 |
-
|
15 |
-
logging.basicConfig(level=logging.INFO)
|
16 |
-
# Constants and Configuration
|
17 |
-
CACHE_DIR = os.getenv("CACHE_DIR")
|
18 |
-
INDEX_FILE = os.getenv("INDEX_FILE")
|
19 |
-
TOKEN = os.getenv("TOKEN")
|
20 |
-
REPO = os.getenv("REPO")
|
21 |
-
|
22 |
-
load_balancer = LoadBalancer(cache_dir=CACHE_DIR, index_file=INDEX_FILE, token=TOKEN, repo=REPO)
|
23 |
-
|
24 |
-
# Start polling in a separate thread
|
25 |
-
polling_thread = Thread(target=load_balancer.start_polling)
|
26 |
-
polling_thread.start()
|
27 |
-
|
28 |
-
# API Endpoints
|
29 |
-
@app.route('/api/film/<title>', methods=['GET'])
|
30 |
-
def get_movie_api(title):
|
31 |
-
"""Endpoint to get the movie by title."""
|
32 |
-
if not title:
|
33 |
-
return jsonify({"error": "Title parameter is required"}), 400
|
34 |
-
|
35 |
-
# Check if the film is already cached
|
36 |
-
if title in load_balancer.FILM_STORE:
|
37 |
-
url = load_balancer.FILM_STORE[title]
|
38 |
-
return jsonify({"url":url})
|
39 |
-
|
40 |
-
movie_path = load_balancer.find_movie_path(title)
|
41 |
-
|
42 |
-
if not movie_path:
|
43 |
-
return jsonify({"error": "Movie not found"}), 404
|
44 |
-
|
45 |
-
# Start the download in a instance
|
46 |
-
response = load_balancer.download_film_to_best_instance(title=title)
|
47 |
-
if response:
|
48 |
-
return jsonify(response)
|
49 |
-
|
50 |
-
@app.route('/api/tv/<title>/<season>/<episode>', methods=['GET'])
|
51 |
-
def get_tv_show_api(title, season, episode):
|
52 |
-
"""Endpoint to get the TV show by title, season, and episode."""
|
53 |
-
if not title or not season or not episode:
|
54 |
-
return jsonify({"error": "Title, season, and episode parameters are required"}), 400
|
55 |
-
|
56 |
-
# Check if the episode is already cached
|
57 |
-
if title in load_balancer.TV_STORE and season in load_balancer.TV_STORE[title]:
|
58 |
-
for ep in load_balancer.TV_STORE[title][season]:
|
59 |
-
if episode in ep:
|
60 |
-
url = load_balancer.TV_STORE[title][season][ep]
|
61 |
-
return jsonify({"url":url})
|
62 |
-
|
63 |
-
tv_path = load_balancer.find_tv_path(title)
|
64 |
-
|
65 |
-
if not tv_path:
|
66 |
-
return jsonify({"error": "TV show not found"}), 404
|
67 |
-
|
68 |
-
episode_path = None
|
69 |
-
for directory in load_balancer.file_structure:
|
70 |
-
if directory['type'] == 'directory' and directory['path'] == 'tv':
|
71 |
-
for sub_directory in directory['contents']:
|
72 |
-
if sub_directory['type'] == 'directory' and title.lower() in sub_directory['path'].lower():
|
73 |
-
for season_dir in sub_directory['contents']:
|
74 |
-
if season_dir['type'] == 'directory' and season in season_dir['path']:
|
75 |
-
for episode_file in season_dir['contents']:
|
76 |
-
if episode_file['type'] == 'file' and episode in episode_file['path']:
|
77 |
-
episode_path = episode_file['path']
|
78 |
-
break
|
79 |
-
|
80 |
-
if not episode_path:
|
81 |
-
return jsonify({"error": "Episode not found"}), 404
|
82 |
-
|
83 |
-
# Start the download in a instance
|
84 |
-
response = load_balancer.download_episode_to_best_instance(title=title, season=season, episode=episode)
|
85 |
-
if response:
|
86 |
-
return jsonify(response)
|
87 |
-
|
88 |
-
@app.route('/api/filmid/<title>', methods=['GET'])
|
89 |
-
def get_film_id_by_title_api(title):
|
90 |
-
"""Endpoint to get the film ID by providing the movie title."""
|
91 |
-
if not title:
|
92 |
-
return jsonify({"error": "Title parameter is required"}), 400
|
93 |
-
film_id = load_balancer.get_film_id(title)
|
94 |
-
return jsonify({"film_id": film_id})
|
95 |
-
|
96 |
-
@app.route('/api/episodeid/<title>/<season>/<episode>', methods=['GET'])
|
97 |
-
def get_episode_id_api(title,season,episode):
|
98 |
-
"""Endpoint to get the episode ID by providing the TV show title, season, and episode."""
|
99 |
-
if not title or not season or not episode:
|
100 |
-
return jsonify({"error": "Title, season, and episode parameters are required"}), 400
|
101 |
-
episode_id = encode_episodeid(title,season,episode)
|
102 |
-
return jsonify({"episode_id": episode_id})
|
103 |
-
|
104 |
-
@app.route('/api/cache/size', methods=['GET'])
|
105 |
-
def get_cache_size_api():
|
106 |
-
total_size = 0
|
107 |
-
for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
|
108 |
-
for f in filenames:
|
109 |
-
fp = os.path.join(dirpath, f)
|
110 |
-
total_size += os.path.getsize(fp)
|
111 |
-
readable_size = bytes_to_human_readable(total_size)
|
112 |
-
return jsonify({"cache_size": readable_size})
|
113 |
-
|
114 |
-
@app.route('/api/cache/clear', methods=['POST'])
|
115 |
-
def clear_cache_api():
|
116 |
-
for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
|
117 |
-
for f in filenames:
|
118 |
-
fp = os.path.join(dirpath, f)
|
119 |
-
os.remove(fp)
|
120 |
-
return jsonify({"status": "Cache cleared"})
|
121 |
-
|
122 |
-
@app.route('/api/tv/store', methods=['GET'])
|
123 |
-
def get_tv_store_api():
|
124 |
-
"""Endpoint to get the TV store JSON."""
|
125 |
-
return jsonify(load_balancer.TV_STORE)
|
126 |
-
|
127 |
-
|
128 |
-
@app.route('/api/film/store', methods=['GET'])
|
129 |
-
def get_film_store_api():
|
130 |
-
"""Endpoint to get the film store JSON."""
|
131 |
-
return jsonify(load_balancer.FILM_STORE)
|
132 |
-
|
133 |
-
|
134 |
-
@app.route('/api/film/metadata/<title>', methods=['GET'])
|
135 |
-
def get_film_metadata_api(title):
|
136 |
-
"""Endpoint to get the film metadata by title."""
|
137 |
-
if not title:
|
138 |
-
return jsonify({'error': 'No title provided'}), 400
|
139 |
-
|
140 |
-
json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
|
141 |
-
|
142 |
-
if os.path.exists(json_cache_path):
|
143 |
-
with open(json_cache_path, 'r') as f:
|
144 |
-
data = json.load(f)
|
145 |
-
return jsonify(data)
|
146 |
-
|
147 |
-
return jsonify({'error': 'Metadata not found'}), 404
|
148 |
-
|
149 |
-
@app.route('/api/tv/metadata/<title>', methods=['GET'])
|
150 |
-
def get_tv_metadata_api(title):
|
151 |
-
"""Endpoint to get the TV show metadata by title."""
|
152 |
-
if not title:
|
153 |
-
return jsonify({'error': 'No title provided'}), 400
|
154 |
-
|
155 |
-
json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
|
156 |
-
|
157 |
-
if os.path.exists(json_cache_path):
|
158 |
-
with open(json_cache_path, 'r') as f:
|
159 |
-
data = json.load(f)
|
160 |
-
|
161 |
-
# Add the file structure to the metadata
|
162 |
-
tv_structure_data = load_balancer.get_tv_structure(title)
|
163 |
-
if tv_structure_data:
|
164 |
-
data['file_structure'] = tv_structure_data
|
165 |
-
|
166 |
-
return jsonify(data)
|
167 |
-
|
168 |
-
return jsonify({'error': 'Metadata not found'}), 404
|
169 |
-
|
170 |
-
|
171 |
-
@app.route("/api/film/all")
|
172 |
-
def get_all_films_api():
|
173 |
-
return load_balancer.get_all_films()
|
174 |
-
|
175 |
-
@app.route("/api/tv/all")
|
176 |
-
def get_all_tvshows_api():
|
177 |
-
return load_balancer.get_all_tv_shows()
|
178 |
-
|
179 |
-
@app.route('/api/instances',methods=["GET"])
|
180 |
-
def get_instances():
|
181 |
-
return load_balancer.instances
|
182 |
-
|
183 |
-
@app.route('/api/instances/health',methods=["GET"])
|
184 |
-
def get_instances_health():
|
185 |
-
return load_balancer.instances_health
|
186 |
-
#############################################################
|
187 |
-
# This API is only for instances
|
188 |
-
@app.route('/api/register', methods=['POST'])
|
189 |
-
def register_instance():
|
190 |
-
try:
|
191 |
-
data = request.json
|
192 |
-
if not data or "url" not in data:
|
193 |
-
return jsonify({"error": "No URL provided"}), 400
|
194 |
-
|
195 |
-
url = data["url"]
|
196 |
-
if not is_valid_url(url):
|
197 |
-
return jsonify({"error": "Invalid URL"}), 400
|
198 |
-
|
199 |
-
# Register the instance
|
200 |
-
load_balancer.register_instance(url)
|
201 |
-
logging.info(f"Instance registered: {url}")
|
202 |
-
|
203 |
-
return jsonify({"message": f"Instance {url} registered successfully"}), 200
|
204 |
-
|
205 |
-
except Exception as e:
|
206 |
-
logging.error(f"Error registering instance: {e}")
|
207 |
-
return jsonify({"error": "Failed to register instance"}), 500
|
208 |
-
|
209 |
-
#############################################################
|
210 |
-
# Routes
|
211 |
-
@app.route('/')
|
212 |
-
def index():
|
213 |
-
return f"Load Balancer is Running {load_balancer.version}"
|
214 |
-
|
215 |
-
# Main entry point
|
216 |
if __name__ == "__main__":
|
217 |
-
|
|
|
|
|
|
|
|
|
|
1 |
+
import subprocess
|
2 |
+
import sys
|
3 |
+
|
4 |
+
def run_uvicorn():
|
5 |
+
# Command to run the FastAPI app with uvicorn
|
6 |
+
command = [
|
7 |
+
"uvicorn",
|
8 |
+
"main:app",
|
9 |
+
"--host", "0.0.0.0",
|
10 |
+
"--port", "7860",
|
11 |
+
"--reload"
|
12 |
+
]
|
13 |
+
|
14 |
+
# Run the command
|
15 |
+
subprocess.run(command, check=True)
|
16 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
17 |
if __name__ == "__main__":
|
18 |
+
try:
|
19 |
+
run_uvicorn()
|
20 |
+
except subprocess.CalledProcessError as e:
|
21 |
+
print(f"Error running uvicorn: {e}", file=sys.stderr)
|
22 |
+
sys.exit(1)
|
main.py
ADDED
@@ -0,0 +1,219 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI, HTTPException, Request
|
2 |
+
from fastapi.responses import JSONResponse
|
3 |
+
from fastapi.middleware.cors import CORSMiddleware
|
4 |
+
from utils import is_valid_url, bytes_to_human_readable, encode_episodeid
|
5 |
+
import os
|
6 |
+
import json
|
7 |
+
import urllib.parse
|
8 |
+
from threading import Thread
|
9 |
+
from LoadBalancer import LoadBalancer
|
10 |
+
import logging
|
11 |
+
|
12 |
+
app = FastAPI()
|
13 |
+
|
14 |
+
# CORS
|
15 |
+
app.add_middleware(
|
16 |
+
CORSMiddleware,
|
17 |
+
allow_origins=["*"],
|
18 |
+
allow_credentials=True,
|
19 |
+
allow_methods=["*"],
|
20 |
+
allow_headers=["*"],
|
21 |
+
)
|
22 |
+
|
23 |
+
logging.basicConfig(level=logging.INFO)
|
24 |
+
|
25 |
+
# Constants and Configuration
|
26 |
+
CACHE_DIR = os.getenv("CACHE_DIR")
|
27 |
+
INDEX_FILE = os.getenv("INDEX_FILE")
|
28 |
+
TOKEN = os.getenv("TOKEN")
|
29 |
+
REPO = os.getenv("REPO")
|
30 |
+
|
31 |
+
load_balancer = LoadBalancer(cache_dir=CACHE_DIR, index_file=INDEX_FILE, token=TOKEN, repo=REPO)
|
32 |
+
|
33 |
+
# Start polling in a separate thread
|
34 |
+
polling_thread = Thread(target=load_balancer.start_polling)
|
35 |
+
polling_thread.start()
|
36 |
+
|
37 |
+
# API Endpoints
|
38 |
+
@app.get("/api/film/{title}")
|
39 |
+
async def get_movie_api(title: str):
|
40 |
+
"""Endpoint to get the movie by title."""
|
41 |
+
if not title:
|
42 |
+
raise HTTPException(status_code=400, detail="Title parameter is required")
|
43 |
+
|
44 |
+
# Check if the film is already cached
|
45 |
+
if title in load_balancer.FILM_STORE:
|
46 |
+
url = load_balancer.FILM_STORE[title]
|
47 |
+
return JSONResponse(content={"url": url})
|
48 |
+
|
49 |
+
movie_path = await load_balancer.find_movie_path(title)
|
50 |
+
|
51 |
+
if not movie_path:
|
52 |
+
raise HTTPException(status_code=404, detail="Movie not found")
|
53 |
+
|
54 |
+
# Start the download in an instance
|
55 |
+
response = await load_balancer.download_film_to_best_instance(title=title)
|
56 |
+
if response:
|
57 |
+
return JSONResponse(content=response)
|
58 |
+
|
59 |
+
@app.get("/api/tv/{title}/{season}/{episode}")
|
60 |
+
async def get_tv_show_api(title: str, season: str, episode: str):
|
61 |
+
"""Endpoint to get the TV show by title, season, and episode."""
|
62 |
+
if not title or not season or not episode:
|
63 |
+
raise HTTPException(status_code=400, detail="Title, season, and episode parameters are required")
|
64 |
+
|
65 |
+
# Check if the episode is already cached
|
66 |
+
if title in load_balancer.TV_STORE and season in load_balancer.TV_STORE[title]:
|
67 |
+
for ep in load_balancer.TV_STORE[title][season]:
|
68 |
+
if episode in ep:
|
69 |
+
url = load_balancer.TV_STORE[title][season][ep]
|
70 |
+
return JSONResponse(content={"url": url})
|
71 |
+
|
72 |
+
tv_path = await load_balancer.find_tv_path(title)
|
73 |
+
|
74 |
+
if not tv_path:
|
75 |
+
raise HTTPException(status_code=404, detail="TV show not found")
|
76 |
+
|
77 |
+
episode_path = None
|
78 |
+
for directory in load_balancer.file_structure:
|
79 |
+
if directory['type'] == 'directory' and directory['path'] == 'tv':
|
80 |
+
for sub_directory in directory['contents']:
|
81 |
+
if sub_directory['type'] == 'directory' and title.lower() in sub_directory['path'].lower():
|
82 |
+
for season_dir in sub_directory['contents']:
|
83 |
+
if season_dir['type'] == 'directory' and season in season_dir['path']:
|
84 |
+
for episode_file in season_dir['contents']:
|
85 |
+
if episode_file['type'] == 'file' and episode in episode_file['path']:
|
86 |
+
episode_path = episode_file['path']
|
87 |
+
break
|
88 |
+
|
89 |
+
if not episode_path:
|
90 |
+
raise HTTPException(status_code=404, detail="Episode not found")
|
91 |
+
|
92 |
+
# Start the download in an instance
|
93 |
+
response = await load_balancer.download_episode_to_best_instance(title=title, season=season, episode=episode)
|
94 |
+
if response:
|
95 |
+
return JSONResponse(content=response)
|
96 |
+
|
97 |
+
@app.get("/api/filmid/{title}")
|
98 |
+
async def get_film_id_by_title_api(title: str):
|
99 |
+
"""Endpoint to get the film ID by providing the movie title."""
|
100 |
+
if not title:
|
101 |
+
raise HTTPException(status_code=400, detail="Title parameter is required")
|
102 |
+
film_id = await load_balancer.get_film_id(title)
|
103 |
+
return JSONResponse(content={"film_id": film_id})
|
104 |
+
|
105 |
+
@app.get("/api/episodeid/{title}/{season}/{episode}")
|
106 |
+
async def get_episode_id_api(title: str, season: str, episode: str):
|
107 |
+
"""Endpoint to get the episode ID by providing the TV show title, season, and episode."""
|
108 |
+
if not title or not season or not episode:
|
109 |
+
raise HTTPException(status_code=400, detail="Title, season, and episode parameters are required")
|
110 |
+
episode_id = encode_episodeid(title, season, episode)
|
111 |
+
return JSONResponse(content={"episode_id": episode_id})
|
112 |
+
|
113 |
+
@app.get("/api/cache/size")
|
114 |
+
async def get_cache_size_api():
|
115 |
+
total_size = 0
|
116 |
+
for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
|
117 |
+
for f in filenames:
|
118 |
+
fp = os.path.join(dirpath, f)
|
119 |
+
total_size += os.path.getsize(fp)
|
120 |
+
readable_size = bytes_to_human_readable(total_size)
|
121 |
+
return JSONResponse(content={"cache_size": readable_size})
|
122 |
+
|
123 |
+
@app.post("/api/cache/clear")
|
124 |
+
async def clear_cache_api():
|
125 |
+
for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
|
126 |
+
for f in filenames:
|
127 |
+
fp = os.path.join(dirpath, f)
|
128 |
+
os.remove(fp)
|
129 |
+
return JSONResponse(content={"status": "Cache cleared"})
|
130 |
+
|
131 |
+
@app.get("/api/tv/store")
|
132 |
+
async def get_tv_store_api():
|
133 |
+
"""Endpoint to get the TV store JSON."""
|
134 |
+
return JSONResponse(content=load_balancer.TV_STORE)
|
135 |
+
|
136 |
+
|
137 |
+
@app.get("/api/film/store")
|
138 |
+
async def get_film_store_api():
|
139 |
+
"""Endpoint to get the film store JSON."""
|
140 |
+
return JSONResponse(content=load_balancer.FILM_STORE)
|
141 |
+
|
142 |
+
@app.get("/api/film/metadata/{title}")
|
143 |
+
async def get_film_metadata_api(title: str):
|
144 |
+
"""Endpoint to get the film metadata by title."""
|
145 |
+
if not title:
|
146 |
+
raise HTTPException(status_code=400, detail="No title provided")
|
147 |
+
|
148 |
+
json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
|
149 |
+
|
150 |
+
if os.path.exists(json_cache_path):
|
151 |
+
with open(json_cache_path, 'r') as f:
|
152 |
+
data = json.load(f)
|
153 |
+
return JSONResponse(content=data)
|
154 |
+
|
155 |
+
raise HTTPException(status_code=404, detail="Metadata not found")
|
156 |
+
|
157 |
+
@app.get("/api/tv/metadata/{title}")
|
158 |
+
async def get_tv_metadata_api(title: str):
|
159 |
+
"""Endpoint to get the TV show metadata by title."""
|
160 |
+
if not title:
|
161 |
+
raise HTTPException(status_code=400, detail="No title provided")
|
162 |
+
|
163 |
+
json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
|
164 |
+
|
165 |
+
if os.path.exists(json_cache_path):
|
166 |
+
with open(json_cache_path, 'r') as f:
|
167 |
+
data = json.load(f)
|
168 |
+
|
169 |
+
# Add the file structure to the metadata
|
170 |
+
tv_structure_data = await load_balancer.get_tv_structure(title)
|
171 |
+
if tv_structure_data:
|
172 |
+
data['file_structure'] = tv_structure_data
|
173 |
+
|
174 |
+
return JSONResponse(content=data)
|
175 |
+
|
176 |
+
raise HTTPException(status_code=404, detail="Metadata not found")
|
177 |
+
|
178 |
+
@app.get("/api/film/all")
|
179 |
+
async def get_all_films_api():
|
180 |
+
return JSONResponse(content=await load_balancer.get_all_films())
|
181 |
+
|
182 |
+
@app.get("/api/tv/all")
|
183 |
+
async def get_all_tvshows_api():
|
184 |
+
return JSONResponse(content=await load_balancer.get_all_tv_shows())
|
185 |
+
|
186 |
+
@app.get("/api/instances")
|
187 |
+
async def get_instances():
|
188 |
+
return JSONResponse(content=load_balancer.instances)
|
189 |
+
|
190 |
+
@app.get("/api/instances/health")
|
191 |
+
async def get_instances_health():
|
192 |
+
return JSONResponse(content=load_balancer.instances_health)
|
193 |
+
|
194 |
+
# This API is only for instances
|
195 |
+
@app.post("/api/register")
|
196 |
+
async def register_instance(request: Request):
|
197 |
+
try:
|
198 |
+
data = await request.json()
|
199 |
+
if not data or "url" not in data:
|
200 |
+
raise HTTPException(status_code=400, detail="No URL provided")
|
201 |
+
|
202 |
+
url = data["url"]
|
203 |
+
if not is_valid_url(url):
|
204 |
+
raise HTTPException(status_code=400, detail="Invalid URL")
|
205 |
+
|
206 |
+
# Register the instance
|
207 |
+
load_balancer.register_instance(url)
|
208 |
+
logging.info(f"Instance registered: {url}")
|
209 |
+
|
210 |
+
return JSONResponse(content={"message": f"Instance {url} registered successfully"}, status_code=200)
|
211 |
+
|
212 |
+
except Exception as e:
|
213 |
+
logging.error(f"Error registering instance: {e}")
|
214 |
+
raise HTTPException(status_code=500, detail="Failed to register instance")
|
215 |
+
|
216 |
+
# Routes
|
217 |
+
@app.get("/")
|
218 |
+
async def index():
|
219 |
+
return f"Load Balancer is Running {load_balancer.version}"
|
requirements.txt
CHANGED
@@ -1,5 +1,6 @@
|
|
1 |
-
|
2 |
-
|
|
|
3 |
requests
|
4 |
python-dotenv
|
5 |
tqdm
|
|
|
1 |
+
fastapi
|
2 |
+
aiofiles
|
3 |
+
uvicorn
|
4 |
requests
|
5 |
python-dotenv
|
6 |
tqdm
|