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
|
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import http.server
import os
import socket
import socketserver
import threading
from pathlib import Path
THREADS = []
class CustomHTTPRequestHandler(http.server.SimpleHTTPRequestHandler):
hub_root = ""
def translate_path(self, path):
# Remove front slash and query args to match the files
return str(self.hub_root / Path(path.lstrip("/").split("?")[0]))
def send_head(self):
path = Path(self.translate_path(self.path))
if path.is_dir():
return super().send_head()
# when dealing with a file, we set the ETag header using the file size.
if path.is_file():
file_size = path.stat().st_size
etag = f'"{file_size}"'
# Handle conditional GET requests
if_match = self.headers.get("If-None-Match")
if if_match == etag:
self.send_response(304)
self.end_headers()
return None
self.send_response(200)
self.send_header("Content-type", self.guess_type(str(path)))
self.send_header("Content-Length", str(file_size))
self.send_header("ETag", etag)
self.end_headers()
return path.open("rb")
self.send_error(404, "File not found")
def serve_directory(directory, port):
"""Serves the directory at the given port."""
CustomHTTPRequestHandler.hub_root = directory
with socketserver.TCPServer(("", port), CustomHTTPRequestHandler) as httpd:
print(f"Serving {directory} at http://localhost:{port}")
httpd.serve_forever()
def start_hub(root_directory):
"""Starts a local hub server and returns the port and thread."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
port = s.getsockname()[1]
server_thread = threading.Thread(
target=serve_directory, args=(root_directory, port), daemon=True
)
server_thread.start()
return port, server_thread
def before_runs(env):
"""Runs before all performance tests.
We grab MOZ_ML_LOCAL_DIR. If set we serve MOZ_ML_LOCAL_DIR/onnx-models as our local hub.
MOZ_FETCHES_DIR is used in the CI as an alternate localtion.
"""
fetches_dir = os.environ.get("MOZ_ML_LOCAL_DIR")
if fetches_dir is None:
fetches_dir = os.environ.get("MOZ_FETCHES_DIR")
if fetches_dir is None:
return
hub_dir = Path(fetches_dir) / "onnx-models"
if not hub_dir.is_dir():
return
port, server_thread = start_hub(hub_dir)
os.environ["MOZ_MODELS_HUB"] = f"http://localhost:{port}"
THREADS.append(server_thread)
def after_runs(env):
if len(THREADS) > 0:
print("Shutting down")
THREADS[0].join(timeout=0)
THREADS.clear()
|
