1
#!/usr/bin/env python3
2
"""Wyoming lazy proxy for faster-whisper.
3

4
Home Assistant talks to this proxy instead of faster-whisper directly.
5
- `describe` events are answered from a cached `info` response, so HA's
6
  availability checks never touch the backend.
7
- The first *real* event of a connection (e.g. `transcribe`) starts the
8
  faster-whisper container via the docker socket, replays the event, and
9
  relays bytes in both directions from then on.
10
- After IDLE_SECONDS without an active session the container is stopped
11
  again, freeing its VRAM.
12

13
stdlib only; runs on any python:3.12 image.
14
"""
15
import asyncio
16
import json
17
import os
18
import socket
19
import time
20

21
BACKEND_HOST = os.environ.get('BACKEND_HOST', '10.0.0.4')
22
BACKEND_PORT = int(os.environ.get('BACKEND_PORT', '10302'))
23
PROXY_PORT = int(os.environ.get('PROXY_PORT', '10300'))
24
IDLE_SECONDS = int(os.environ.get('IDLE_SECONDS', '900'))
25
CONTAINER = os.environ.get('CONTAINER', 'faster-whisper')
26
DOCKER_SOCK = os.environ.get('DOCKER_SOCK', '/var/run/docker.sock')
27
CACHE_FILE = os.environ.get('CACHE_FILE', '/app/info.cache')
28
START_TIMEOUT = int(os.environ.get('START_TIMEOUT', '180'))
29

30
state = {'active': 0, 'last': time.monotonic(), 'lock': None}
31

32

33
def log(*args):
34
    print(time.strftime('%H:%M:%S'), *args, flush=True)
35

36

37
def docker_api(method, path):
38
    """Minimal docker engine API call over the unix socket. Returns (status, body)."""
39
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
40
    sock.settimeout(60)
41
    sock.connect(DOCKER_SOCK)
42
    sock.sendall(f'{method} {path} HTTP/1.1\r\nHost: docker\r\nConnection: close\r\n\r\n'.encode())
43
    buf = b''
44
    while True:
45
        chunk = sock.recv(65536)
46
        if not chunk:
47
            break
48
        buf += chunk
49
    sock.close()
50
    head, _, body = buf.partition(b'\r\n\r\n')
51
    return int(head.split(b' ', 2)[1]), body
52

53

54
def backend_running():
55
    try:
56
        status, body = docker_api('GET', f'/containers/{CONTAINER}/json')
57
        if status != 200:
58
            return False
59
        text = body.decode('utf-8', 'replace')
60
        text = text[text.find('{'):text.rfind('}') + 1]
61
        return json.loads(text).get('State', {}).get('Running', False)
62
    except Exception as err:
63
        log('backend_running check failed:', err)
64
        return False
65

66

67
def start_container():
68
    status, _ = docker_api('POST', f'/containers/{CONTAINER}/start')
69
    return status in (204, 304)
70

71

72
def stop_container():
73
    status, _ = docker_api('POST', f'/containers/{CONTAINER}/stop?t=30')
74
    return status in (204, 304)
75

76

77
async def read_event(reader):
78
    """Reads one wyoming event; returns {'raw': bytes, 'type': str} or None on EOF."""
79
    line = await reader.readline()
80
    if not line:
81
        return None
82
    header = json.loads(line)
83
    extra = b''
84
    for key in ('data_length', 'payload_length'):
85
        length = int(header.get(key) or 0)
86
        if length:
87
            extra += await reader.readexactly(length)
88
    return {'raw': line + extra, 'type': header.get('type')}
89

90

91
async def ensure_backend():
92
    """Starts the backend container if needed and returns an open (reader, writer).
93

94
    A bare TCP connect is not proof of readiness: docker-proxy accepts on the
95
    published port as soon as the container starts, long before the Wyoming
96
    server listens. So readiness = a real describe -> info handshake.
97
    """
98
    async with state['lock']:
99
        deadline = time.monotonic() + START_TIMEOUT
100
        if not await asyncio.to_thread(backend_running):
101
            log('starting backend container')
102
            await asyncio.to_thread(start_container)
103
        while time.monotonic() < deadline:
104
            try:
105
                reader, writer = await asyncio.open_connection(BACKEND_HOST, BACKEND_PORT)
106
                writer.write(json.dumps({'type': 'describe', 'data': {}}).encode() + b'\n')
107
                await writer.drain()
108
                event = await asyncio.wait_for(read_event(reader), 5)
109
                if event and event['type'] == 'info':
110
                    # Connection is clean again after consuming the info event.
111
                    return reader, writer
112
                writer.close()
113
            except (OSError, asyncio.TimeoutError, asyncio.IncompleteReadError, json.JSONDecodeError):
114
                pass
115
            await asyncio.sleep(1)
116
        raise RuntimeError('backend did not come up within START_TIMEOUT')
117

118

119
async def refresh_cache():
120
    try:
121
        reader, writer = await asyncio.open_connection(BACKEND_HOST, BACKEND_PORT)
122
        writer.write(json.dumps({'type': 'describe', 'data': {}}).encode() + b'\n')
123
        await writer.drain()
124
        event = await asyncio.wait_for(read_event(reader), 20)
125
        writer.close()
126
        if event and event['type'] == 'info':
127
            with open(CACHE_FILE, 'wb') as f:
128
                f.write(event['raw'])
129
            log('info cache refreshed:', len(event['raw']), 'bytes')
130
    except Exception as err:
131
        log('cache refresh failed:', err)
132

133

134
async def pipe(reader, writer):
135
    try:
136
        while True:
137
            chunk = await reader.read(65536)
138
            if not chunk:
139
                break
140
            writer.write(chunk)
141
            await writer.drain()
142
    except (ConnectionResetError, BrokenPipeError):
143
        pass
144
    finally:
145
        try:
146
            writer.close()
147
        except Exception:
148
            pass
149

150

151
async def handle_client(client_reader, client_writer):
152
    peer = client_writer.get_extra_info('peername')
153
    backend = None
154
    try:
155
        while backend is None:
156
            event = await read_event(client_reader)
157
            if event is None:
158
                return
159
            if event['type'] == 'ping':
160
                client_writer.write(json.dumps({'type': 'pong', 'data': {}}).encode() + b'\n')
161
                await client_writer.drain()
162
                continue
163
            if event['type'] == 'describe' and os.path.exists(CACHE_FILE):
164
                with open(CACHE_FILE, 'rb') as f:
165
                    client_writer.write(f.read())
166
                await client_writer.drain()
167
                continue
168
            # Real work (or describe without a cache yet): wake the backend.
169
            state['active'] += 1
170
            state['last'] = time.monotonic()
171
            try:
172
                backend = await ensure_backend()
173
            except Exception as err:
174
                state['active'] -= 1
175
                state['last'] = time.monotonic()
176
                log('backend wake failed:', err)
177
                return
178
            log(f'session start from {peer} (trigger: {event["type"]})')
179
            backend[1].write(event['raw'])
180
            await backend[1].drain()
181
            asyncio.ensure_future(refresh_cache())
182
        await asyncio.gather(pipe(client_reader, backend[1]), pipe(backend[0], client_writer))
183
    except (ConnectionResetError, BrokenPipeError, asyncio.IncompleteReadError, json.JSONDecodeError) as err:
184
        log('client error:', err)
185
    finally:
186
        if backend is not None:
187
            state['active'] -= 1
188
            state['last'] = time.monotonic()
189
            log('session end, active:', state['active'])
190
        try:
191
            client_writer.close()
192
        except Exception:
193
            pass
194

195

196
async def idle_reaper():
197
    while True:
198
        await asyncio.sleep(30)
199
        idle = time.monotonic() - state['last']
200
        if state['active'] == 0 and idle > IDLE_SECONDS and await asyncio.to_thread(backend_running):
201
            log(f'idle for {int(idle)}s - stopping backend container')
202
            await asyncio.to_thread(stop_container)
203

204

205
async def main():
206
    state['lock'] = asyncio.Lock()
207
    if not os.path.exists(CACHE_FILE):
208
        log('no info cache - bootstrapping backend once to capture it')
209
        try:
210
            reader, writer = await ensure_backend()
211
            writer.close()
212
            await refresh_cache()
213
        except Exception as err:
214
            log('bootstrap failed (will retry on first request):', err)
215
    server = await asyncio.start_server(handle_client, '0.0.0.0', PROXY_PORT)
216
    log(f'wyoming lazy proxy listening on :{PROXY_PORT} -> {BACKEND_HOST}:{BACKEND_PORT} '
217
        f'(container {CONTAINER}, idle {IDLE_SECONDS}s)')
218
    async with server:
219
        await asyncio.gather(server.serve_forever(), idle_reaper())
220

221

222
if __name__ == '__main__':
223
    asyncio.run(main())