File size: 12,098 Bytes
9c6594c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
#
"""Internal process.

This module implements the entrypoint for the internal process. The internal process
is responsible for handling "record" requests, and responding with "results". Data is
passed to the process over multiprocessing queues.

Threads:
    HandlerThread -- read from record queue and call handlers
    SenderThread -- send to network
    WriterThread -- write to disk

"""

import atexit
import logging
import os
import queue
import sys
import threading
import time
import traceback
from datetime import datetime
from typing import TYPE_CHECKING, Any, List, Optional

import psutil

import wandb

from ..interface.interface_queue import InterfaceQueue
from . import context, handler, internal_util, sender, writer

if TYPE_CHECKING:
    from queue import Queue
    from threading import Event

    from wandb.proto.wandb_internal_pb2 import Record, Result

    from .internal_util import RecordLoopThread
    from .settings_static import SettingsStatic


logger = logging.getLogger(__name__)


def wandb_internal(
    settings: "SettingsStatic",
    record_q: "Queue[Record]",
    result_q: "Queue[Result]",
    port: Optional[int] = None,
    user_pid: Optional[int] = None,
) -> None:
    """Internal process function entrypoint.

    Read from record queue and dispatch work to various threads.

    Args:
        settings: settings object
        record_q: records to be handled
        result_q: for sending results back

    """
    # mark this process as internal
    wandb._set_internal_process()  # type: ignore
    started = time.time()

    # any sentry events in the internal process will be tagged as such
    wandb._sentry.configure_scope(process_context="internal", tags=dict(settings))

    # register the exit handler only when wandb_internal is called, not on import
    @atexit.register
    def handle_exit(*args: "Any") -> None:
        logger.info("Internal process exited")

    # Let's make sure we don't modify settings so use a static object
    _settings = settings
    if _settings.log_internal:
        configure_logging(_settings.log_internal, _settings.x_log_level)

    user_pid = user_pid or os.getppid()
    pid = os.getpid()

    logger.info(
        "W&B internal server running at pid: %s, started at: %s",
        pid,
        datetime.fromtimestamp(started),
    )

    publish_interface = InterfaceQueue(record_q=record_q)

    stopped = threading.Event()
    threads: List[RecordLoopThread] = []

    context_keeper = context.ContextKeeper()

    send_record_q: Queue[Record] = queue.Queue()

    write_record_q: Queue[Record] = queue.Queue()

    record_sender_thread = SenderThread(
        settings=_settings,
        record_q=send_record_q,
        result_q=result_q,
        stopped=stopped,
        interface=publish_interface,
        debounce_interval_ms=5000,
        context_keeper=context_keeper,
    )
    threads.append(record_sender_thread)

    record_writer_thread = WriterThread(
        settings=_settings,
        record_q=write_record_q,
        result_q=result_q,
        stopped=stopped,
        interface=publish_interface,
        sender_q=send_record_q,
        context_keeper=context_keeper,
    )
    threads.append(record_writer_thread)

    record_handler_thread = HandlerThread(
        settings=_settings,
        record_q=record_q,
        result_q=result_q,
        stopped=stopped,
        writer_q=write_record_q,
        interface=publish_interface,
        context_keeper=context_keeper,
    )
    threads.append(record_handler_thread)

    process_check = ProcessCheck(settings=_settings, user_pid=user_pid)

    for thread in threads:
        thread.start()

    interrupt_count = 0
    while not stopped.is_set():
        try:
            # wait for stop event
            while not stopped.is_set():
                time.sleep(1)
                if process_check.is_dead():
                    logger.error("Internal process shutdown.")
                    stopped.set()
        except KeyboardInterrupt:
            interrupt_count += 1
            logger.warning(f"Internal process interrupt: {interrupt_count}")
        finally:
            if interrupt_count >= 2:
                logger.error("Internal process interrupted.")
                stopped.set()

    for thread in threads:
        thread.join()

    def close_internal_log() -> None:
        root = logging.getLogger("wandb")
        for _handler in root.handlers[:]:
            _handler.close()
            root.removeHandler(_handler)

    for thread in threads:
        exc_info = thread.get_exception()
        if exc_info:
            logger.error(f"Thread {thread.name}:", exc_info=exc_info)
            print(f"Thread {thread.name}:", file=sys.stderr)  # noqa: T201
            traceback.print_exception(*exc_info)
            wandb._sentry.exception(exc_info)
            wandb.termerror("Internal wandb error: file data was not synced")
            # TODO: We can make this more graceful by returning an error to streams.py
            # and potentially just fail the one stream.
            os._exit(-1)

    close_internal_log()


def configure_logging(
    log_fname: str, log_level: int, run_id: Optional[str] = None
) -> None:
    # TODO: we may want make prints and stdout make it into the logs
    # sys.stdout = open(settings.log_internal, "a")
    # sys.stderr = open(settings.log_internal, "a")
    log_handler = logging.FileHandler(log_fname)
    log_handler.setLevel(log_level)

    class WBFilter(logging.Filter):
        def filter(self, record: "Any") -> bool:
            record.run_id = run_id
            return True

    if run_id:
        formatter = logging.Formatter(
            "%(asctime)s %(levelname)-7s %(threadName)-10s:%(process)d "
            "[%(run_id)s:%(filename)s:%(funcName)s():%(lineno)s] %(message)s"
        )
    else:
        formatter = logging.Formatter(
            "%(asctime)s %(levelname)-7s %(threadName)-10s:%(process)d "
            "[%(filename)s:%(funcName)s():%(lineno)s] %(message)s"
        )

    log_handler.setFormatter(formatter)
    if run_id:
        log_handler.addFilter(WBFilter())
    # If this is called without "wandb", backend logs from this module
    # are not streamed to `debug-internal.log` when we spawn with fork
    # TODO: (cvp) we should really take another pass at logging in general
    root = logging.getLogger("wandb")
    root.propagate = False
    root.setLevel(logging.DEBUG)
    root.addHandler(log_handler)


class HandlerThread(internal_util.RecordLoopThread):
    """Read records from queue and dispatch to handler routines."""

    _record_q: "Queue[Record]"
    _result_q: "Queue[Result]"
    _stopped: "Event"
    _context_keeper: context.ContextKeeper

    def __init__(
        self,
        settings: "SettingsStatic",
        record_q: "Queue[Record]",
        result_q: "Queue[Result]",
        stopped: "Event",
        writer_q: "Queue[Record]",
        interface: "InterfaceQueue",
        context_keeper: context.ContextKeeper,
        debounce_interval_ms: "float" = 1000,
    ) -> None:
        super().__init__(
            input_record_q=record_q,
            result_q=result_q,
            stopped=stopped,
            debounce_interval_ms=debounce_interval_ms,
        )
        self.name = "HandlerThread"
        self._settings = settings
        self._record_q = record_q
        self._result_q = result_q
        self._stopped = stopped
        self._writer_q = writer_q
        self._interface = interface
        self._context_keeper = context_keeper

    def _setup(self) -> None:
        self._hm = handler.HandleManager(
            settings=self._settings,
            record_q=self._record_q,
            result_q=self._result_q,
            stopped=self._stopped,
            writer_q=self._writer_q,
            interface=self._interface,
            context_keeper=self._context_keeper,
        )

    def _process(self, record: "Record") -> None:
        self._hm.handle(record)

    def _finish(self) -> None:
        self._hm.finish()

    def _debounce(self) -> None:
        self._hm.debounce()


class SenderThread(internal_util.RecordLoopThread):
    """Read records from queue and dispatch to sender routines."""

    _record_q: "Queue[Record]"
    _result_q: "Queue[Result]"
    _context_keeper: context.ContextKeeper

    def __init__(
        self,
        settings: "SettingsStatic",
        record_q: "Queue[Record]",
        result_q: "Queue[Result]",
        stopped: "Event",
        interface: "InterfaceQueue",
        context_keeper: context.ContextKeeper,
        debounce_interval_ms: "float" = 5000,
    ) -> None:
        super().__init__(
            input_record_q=record_q,
            result_q=result_q,
            stopped=stopped,
            debounce_interval_ms=debounce_interval_ms,
        )
        self.name = "SenderThread"
        self._settings = settings
        self._record_q = record_q
        self._result_q = result_q
        self._interface = interface
        self._context_keeper = context_keeper

    def _setup(self) -> None:
        self._sm = sender.SendManager(
            settings=self._settings,
            record_q=self._record_q,
            result_q=self._result_q,
            interface=self._interface,
            context_keeper=self._context_keeper,
        )

    def _process(self, record: "Record") -> None:
        self._sm.send(record)

    def _finish(self) -> None:
        self._sm.finish()

    def _debounce(self) -> None:
        self._sm.debounce()


class WriterThread(internal_util.RecordLoopThread):
    """Read records from queue and dispatch to writer routines."""

    _record_q: "Queue[Record]"
    _result_q: "Queue[Result]"
    _context_keeper: context.ContextKeeper

    def __init__(
        self,
        settings: "SettingsStatic",
        record_q: "Queue[Record]",
        result_q: "Queue[Result]",
        stopped: "Event",
        interface: "InterfaceQueue",
        sender_q: "Queue[Record]",
        context_keeper: context.ContextKeeper,
        debounce_interval_ms: "float" = 1000,
    ) -> None:
        super().__init__(
            input_record_q=record_q,
            result_q=result_q,
            stopped=stopped,
            debounce_interval_ms=debounce_interval_ms,
        )
        self.name = "WriterThread"
        self._settings = settings
        self._record_q = record_q
        self._result_q = result_q
        self._sender_q = sender_q
        self._interface = interface
        self._context_keeper = context_keeper

    def _setup(self) -> None:
        self._wm = writer.WriteManager(
            settings=self._settings,
            record_q=self._record_q,
            result_q=self._result_q,
            sender_q=self._sender_q,
            interface=self._interface,
            context_keeper=self._context_keeper,
        )

    def _process(self, record: "Record") -> None:
        self._wm.write(record)

    def _finish(self) -> None:
        self._wm.finish()

    def _debounce(self) -> None:
        self._wm.debounce()


class ProcessCheck:
    """Class to help watch a process id to detect when it is dead."""

    check_process_last: Optional[float]

    def __init__(self, settings: "SettingsStatic", user_pid: Optional[int]) -> None:
        self.settings = settings
        self.pid = user_pid
        self.check_process_last = None
        self.check_process_interval = settings.x_internal_check_process

    def is_dead(self) -> bool:
        if not self.check_process_interval or not self.pid:
            return False
        time_now = time.time()
        if (
            self.check_process_last
            and time_now < self.check_process_last + self.check_process_interval
        ):
            return False
        self.check_process_last = time_now

        # TODO(jhr): check for os.getppid on unix being 1?
        exists = psutil.pid_exists(self.pid)
        if not exists:
            logger.warning(
                f"Internal process exiting, parent pid {self.pid} disappeared"
            )
            return True
        return False