File size: 16,693 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
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
"""Global W&B library state.

This module manages global state, which for wandb includes:

- Settings configured through `wandb.setup()`
- The list of active runs
- A subprocess ("the internal service") that asynchronously uploads metrics

This module is fork-aware: in a forked process such as that spawned by the
`multiprocessing` module, `wandb.singleton()` returns a new object, not the
one inherited from the parent process. This requirement comes from backward
compatibility with old design choices: the hardest one to fix is that wandb
was originally designed to have a single run for the entire process that
`wandb.init()` was meant to return. Back then, the only way to create
multiple simultaneous runs in a single script was to run subprocesses, and since
the built-in `multiprocessing` module forks by default, this required a PID
check to make `wandb.init()` ignore the inherited global run.

Another reason for fork-awareness is that the process that starts up
the internal service owns it and is responsible for shutting it down,
and child processes shouldn't also try to do that. This is easier to
redesign.
"""

from __future__ import annotations

import logging
import os
import pathlib
import sys
import threading
from typing import TYPE_CHECKING, Any, Union

import wandb
import wandb.integration.sagemaker as sagemaker
from wandb.env import CONFIG_DIR
from wandb.sdk.lib import import_hooks, wb_logging

from . import wandb_settings
from .lib import config_util, server

if TYPE_CHECKING:
    from wandb.sdk import wandb_run
    from wandb.sdk.lib.service_connection import ServiceConnection
    from wandb.sdk.wandb_settings import Settings


class _EarlyLogger:
    """Early logger which captures logs in memory until logging can be configured."""

    def __init__(self) -> None:
        self._log: list[tuple] = []
        self._exception: list[tuple] = []
        # support old warn() as alias of warning()
        self.warn = self.warning

    def debug(self, msg: str, *args: Any, **kwargs: Any) -> None:
        self._log.append((logging.DEBUG, msg, args, kwargs))

    def info(self, msg: str, *args: Any, **kwargs: Any) -> None:
        self._log.append((logging.INFO, msg, args, kwargs))

    def warning(self, msg: str, *args: Any, **kwargs: Any) -> None:
        self._log.append((logging.WARNING, msg, args, kwargs))

    def error(self, msg: str, *args: Any, **kwargs: Any) -> None:
        self._log.append((logging.ERROR, msg, args, kwargs))

    def critical(self, msg: str, *args: Any, **kwargs: Any) -> None:
        self._log.append((logging.CRITICAL, msg, args, kwargs))

    def exception(self, msg: str, *args: Any, **kwargs: Any) -> None:
        self._exception.append((msg, args, kwargs))

    def log(self, level: str, msg: str, *args: Any, **kwargs: Any) -> None:
        self._log.append((level, msg, args, kwargs))

    def _flush(self, new_logger: Logger) -> None:
        assert self is not new_logger
        for level, msg, args, kwargs in self._log:
            new_logger.log(level, msg, *args, **kwargs)
        for msg, args, kwargs in self._exception:
            new_logger.exception(msg, *args, **kwargs)


Logger = Union[logging.Logger, _EarlyLogger]


class _WandbSetup:
    """W&B library singleton."""

    def __init__(
        self,
        pid: int,
        settings: Settings | None = None,
        environ: dict | None = None,
    ) -> None:
        self._connection: ServiceConnection | None = None

        self._active_runs: list[wandb_run.Run] = []

        self._environ = environ or dict(os.environ)
        self._sweep_config: dict | None = None
        self._config: dict | None = None
        self._server: server.Server | None = None
        self._pid = pid

        # TODO(jhr): defer strict checks until settings are fully initialized
        #            and logging is ready
        self._logger: Logger = _EarlyLogger()

        self._settings = self._settings_setup(settings)

        wandb.termsetup(self._settings, None)

        self._check()
        self._setup()

    def add_active_run(self, run: wandb_run.Run) -> None:
        """Append a run to the active runs list.

        This must be called when a run is initialized.

        Args:
            run: A newly initialized run.
        """
        if run not in self._active_runs:
            self._active_runs.append(run)

    def remove_active_run(self, run: wandb_run.Run) -> None:
        """Remove the run from the active runs list.

        This must be called when a run is finished.

        Args:
            run: A run that is finished or crashed.
        """
        try:
            self._active_runs.remove(run)
        except ValueError:
            pass  # Removing a run multiple times is not an error.

    @property
    def most_recent_active_run(self) -> wandb_run.Run | None:
        """The most recently initialized run that is not yet finished."""
        if not self._active_runs:
            return None

        return self._active_runs[-1]

    def finish_all_active_runs(self) -> None:
        """Finish all unfinished runs.

        NOTE: This is slightly inefficient as it finishes runs one at a time.
        This only exists to support using the `reinit="finish_previous"`
        setting together with `reinit="create_new"` which does not seem to be a
        useful pattern. Since `"create_new"` should eventually become the
        default and only behavior, it does not seem worth optimizing.
        """
        # Take a snapshot as each call to `finish()` modifies `_active_runs`.
        runs_copy = list(self._active_runs)
        for run in runs_copy:
            run.finish()

    def _settings_setup(
        self,
        settings: Settings | None,
    ) -> wandb_settings.Settings:
        s = wandb_settings.Settings()

        # the pid of the process to monitor for system stats
        pid = os.getpid()
        self._logger.info(f"Current SDK version is {wandb.__version__}")
        self._logger.info(f"Configure stats pid to {pid}")
        s.x_stats_pid = pid

        if settings and settings.settings_system:
            s.settings_system = settings.settings_system
        elif config_dir_str := os.getenv(CONFIG_DIR, None):
            config_dir = pathlib.Path(config_dir_str).expanduser()
            s.settings_system = str(config_dir / "settings")
        else:
            s.settings_system = str(
                pathlib.Path("~", ".config", "wandb", "settings").expanduser()
            )

        # load settings from the system config
        if s.settings_system:
            self._logger.info(f"Loading settings from {s.settings_system}")
        s.update_from_system_config_file()

        # load settings from the workspace config
        if s.settings_workspace:
            self._logger.info(f"Loading settings from {s.settings_workspace}")
        s.update_from_workspace_config_file()

        # load settings from the environment variables
        self._logger.info("Loading settings from environment variables")
        s.update_from_env_vars(self._environ)

        # infer settings from the system environment
        s.update_from_system_environment()

        # load SageMaker settings
        check_sagemaker_env = not s.sagemaker_disable
        if settings and settings.sagemaker_disable:
            check_sagemaker_env = False
        if check_sagemaker_env and sagemaker.is_using_sagemaker():
            self._logger.info("Loading SageMaker settings")
            sagemaker.set_global_settings(s)

        # load settings from the passed init/setup settings
        if settings:
            s.update_from_settings(settings)

        return s

    def _update(self, settings: Settings | None = None) -> None:
        if not settings:
            return
        self._settings.update_from_settings(settings)

    def _update_user_settings(self) -> None:
        # Get rid of cached results to force a refresh.
        self._server = None
        user_settings = self._load_user_settings()
        if user_settings is not None:
            self._settings.update_from_dict(user_settings)

    def _early_logger_flush(self, new_logger: Logger) -> None:
        if self._logger is new_logger:
            return

        if isinstance(self._logger, _EarlyLogger):
            self._logger._flush(new_logger)
        self._logger = new_logger

    def _get_logger(self) -> Logger:
        return self._logger

    @property
    def settings(self) -> wandb_settings.Settings:
        return self._settings

    def _get_entity(self) -> str | None:
        if self._settings and self._settings._offline:
            return None
        entity = self.viewer.get("entity")
        return entity

    def _get_username(self) -> str | None:
        if self._settings and self._settings._offline:
            return None
        return self.viewer.get("username")

    def _get_teams(self) -> list[str]:
        if self._settings and self._settings._offline:
            return []
        teams = self.viewer.get("teams")
        if teams:
            teams = [team["node"]["name"] for team in teams["edges"]]
        return teams or []

    @property
    def viewer(self) -> dict[str, Any]:
        if self._server is None:
            self._server = server.Server(settings=self._settings)

        return self._server.viewer

    def _load_user_settings(self) -> dict[str, Any] | None:
        # offline?
        if self._server is None:
            return None

        flags = self._server._flags
        user_settings = dict()
        if "code_saving_enabled" in flags:
            user_settings["save_code"] = flags["code_saving_enabled"]

        email = self.viewer.get("email", None)
        if email:
            user_settings["email"] = email

        return user_settings

    def _check(self) -> None:
        if hasattr(threading, "main_thread"):
            if threading.current_thread() is not threading.main_thread():
                pass
        elif threading.current_thread().name != "MainThread":
            wandb.termwarn(f"bad thread2: {threading.current_thread().name}")
        if getattr(sys, "frozen", False):
            wandb.termwarn("frozen, could be trouble")

    def _setup(self) -> None:
        sweep_path = self._settings.sweep_param_path
        if sweep_path:
            self._sweep_config = config_util.dict_from_config_file(
                sweep_path, must_exist=True
            )

        # if config_paths was set, read in config dict
        if self._settings.config_paths:
            # TODO(jhr): handle load errors, handle list of files
            for config_path in self._settings.config_paths:
                config_dict = config_util.dict_from_config_file(config_path)
                if config_dict is None:
                    continue
                if self._config is not None:
                    self._config.update(config_dict)
                else:
                    self._config = config_dict

    def _teardown(self, exit_code: int | None = None) -> None:
        import_hooks.unregister_all_post_import_hooks()

        if not self._connection:
            return

        # Reset to None so that setup() creates a new connection.
        connection = self._connection
        self._connection = None

        internal_exit_code = connection.teardown(exit_code or 0)
        if internal_exit_code not in (None, 0):
            sys.exit(internal_exit_code)

    def ensure_service(self) -> ServiceConnection:
        """Returns a connection to the service process creating it if needed."""
        if self._connection:
            return self._connection

        from wandb.sdk.lib import service_connection

        self._connection = service_connection.connect_to_service(self._settings)
        return self._connection

    def assert_service(self) -> ServiceConnection:
        """Returns a connection to the service process, asserting it exists.

        Unlike ensure_service(), this will not start up a service process
        if it didn't already exist.
        """
        if not self._connection:
            raise AssertionError("Expected service process to exist.")

        return self._connection


_singleton: _WandbSetup | None = None
"""The W&B library singleton, or None if not yet set up.

The value is invalid and must not be used if `os.getpid() != _singleton._pid`.
"""


def singleton() -> _WandbSetup:
    """The W&B singleton for the current process.

    The first call to this in this process (which may be a fork of another
    process) creates the singleton, and all subsequent calls return it
    until teardown(). This does not start the service process.
    """
    return _setup(start_service=False)


def singleton_if_setup() -> _WandbSetup | None:
    """The W&B singleton for the current process or None if it isn't set up.

    Always prefer singleton() over this function.

    Unlike singleton(), this never creates the singleton and therefore never
    initializes global settings from the environment. This is useful only
    during tests, which may modify the environment after having imported wandb
    and called certain functions.
    """
    if _singleton and _singleton._pid == os.getpid():
        return _singleton
    else:
        return None


@wb_logging.log_to_all_runs()
def _setup(
    settings: Settings | None = None,
    start_service: bool = True,
) -> _WandbSetup:
    """Set up library context.

    Args:
        settings: Global settings to set, or updates to the global settings
            if the singleton has already been initialized.
        start_service: Whether to start up the service process.
            NOTE: A service process will only be started if allowed by the
            global settings (after the given updates). The service will not
            start up if the mode resolves to "disabled".
    """
    global _singleton

    pid = os.getpid()

    if _singleton and _singleton._pid == pid:
        _singleton._update(settings=settings)
    else:
        _singleton = _WandbSetup(settings=settings, pid=pid)

    if start_service and not _singleton.settings._noop:
        _singleton.ensure_service()

    return _singleton


def setup(settings: Settings | None = None) -> _WandbSetup:
    """Prepares W&B for use in the current process and its children.

    You can usually ignore this as it is implicitly called by `wandb.init()`.

    When using wandb in multiple processes, calling `wandb.setup()`
    in the parent process before starting child processes may improve
    performance and resource utilization.

    Note that `wandb.setup()` modifies `os.environ`, and it is important
    that child processes inherit the modified environment variables.

    See also `wandb.teardown()`.

    Args:
        settings: Configuration settings to apply globally. These can be
            overridden by subsequent `wandb.init()` calls.

    Example:
        ```python
        import multiprocessing

        import wandb


        def run_experiment(params):
            with wandb.init(config=params):
                # Run experiment
                pass


        if __name__ == "__main__":
            # Start backend and set global config
            wandb.setup(settings={"project": "my_project"})

            # Define experiment parameters
            experiment_params = [
                {"learning_rate": 0.01, "epochs": 10},
                {"learning_rate": 0.001, "epochs": 20},
            ]

            # Start multiple processes, each running a separate experiment
            processes = []
            for params in experiment_params:
                p = multiprocessing.Process(target=run_experiment, args=(params,))
                p.start()
                processes.append(p)

            # Wait for all processes to complete
            for p in processes:
                p.join()

            # Optional: Explicitly shut down the backend
            wandb.teardown()
        ```
    """
    return _setup(settings=settings)


@wb_logging.log_to_all_runs()
def teardown(exit_code: int | None = None) -> None:
    """Waits for wandb to finish and frees resources.

    Completes any runs that were not explicitly finished
    using `run.finish()` and waits for all data to be uploaded.

    It is recommended to call this at the end of a session
    that used `wandb.setup()`. It is invoked automatically
    in an `atexit` hook, but this is not reliable in certain setups
    such as when using Python's `multiprocessing` module.
    """
    global _singleton

    orig_singleton = _singleton
    _singleton = None

    if orig_singleton:
        orig_singleton._teardown(exit_code=exit_code)