File size: 11,815 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
import asyncio
import logging
import os
import sys
from typing import Any, Dict, List, Optional, Tuple

import yaml

import wandb
from wandb.apis.internal import Api

from . import loader
from ._project_spec import LaunchProject
from .agent import LaunchAgent
from .agent.agent import construct_agent_configs
from .environment.local_environment import LocalEnvironment
from .errors import ExecutionError, LaunchError
from .runner.abstract import AbstractRun
from .utils import (
    LAUNCH_CONFIG_FILE,
    PROJECT_SYNCHRONOUS,
    construct_launch_spec,
    validate_launch_spec_source,
)

_logger = logging.getLogger(__name__)


def set_launch_logfile(logfile: str) -> None:
    """Set the logfile for the launch agent."""
    # Get logger of parent module
    _launch_logger = logging.getLogger("wandb.sdk.launch")
    if logfile == "-":
        logfile_stream = sys.stdout
    else:
        try:
            logfile_stream = open(logfile, "w")
        # check if file is writable
        except Exception as e:
            wandb.termerror(
                f"Could not open {logfile} for writing logs. Please check "
                f"the path and permissions.\nError: {e}"
            )
            return

    wandb.termlog(
        f"Internal agent logs printing to {'stdout' if logfile == '-' else logfile}. "
    )
    handler = logging.StreamHandler(logfile_stream)
    handler.formatter = logging.Formatter(
        "%(asctime)s %(levelname)-7s %(threadName)-10s:%(process)d "
        "[%(filename)s:%(funcName)s():%(lineno)s] %(message)s"
    )
    _launch_logger.addHandler(handler)
    _launch_logger.log(logging.INFO, "Internal agent logs printing to %s", logfile)


def resolve_agent_config(
    entity: Optional[str],
    max_jobs: Optional[int],
    queues: Optional[Tuple[str]],
    config: Optional[str],
    verbosity: Optional[int],
) -> Tuple[Dict[str, Any], Api]:
    """Resolve the agent config.

    Arguments:
        api (Api): The api.
        entity (str): The entity.
        max_jobs (int): The max number of jobs.
        queues (Tuple[str]): The queues.
        config (str): The config.
        verbosity (int): How verbose to print, 0 or None = default, 1 = print status every 20 seconds, 2 = also print debugging information

    Returns:
        Tuple[Dict[str, Any], Api]: The resolved config and api.
    """
    defaults = {
        "max_jobs": 1,
        "max_schedulers": 1,
        "queues": [],
        "registry": {},
        "builder": {},
        "verbosity": 0,
    }
    resolved_config: Dict[str, Any] = defaults
    config_path = config or os.path.expanduser(LAUNCH_CONFIG_FILE)
    if os.path.isfile(config_path):
        launch_config = {}
        with open(config_path) as f:
            try:
                launch_config = yaml.safe_load(f)
                # This is considered unreachable by mypy, but it's not.
                if launch_config is None:
                    launch_config = {}  # type: ignore
            except yaml.YAMLError as e:
                raise LaunchError(f"Invalid launch agent config: {e}")
        resolved_config.update(launch_config.items())
    elif config is not None:
        raise LaunchError(
            f"Could not find use specified launch config file: {config_path}"
        )
    if os.environ.get("WANDB_ENTITY") is not None:
        resolved_config.update({"entity": os.environ.get("WANDB_ENTITY")})
    if os.environ.get("WANDB_LAUNCH_MAX_JOBS") is not None:
        resolved_config.update(
            {"max_jobs": int(os.environ.get("WANDB_LAUNCH_MAX_JOBS", 1))}
        )

    if entity is not None:
        resolved_config.update({"entity": entity})
    if max_jobs is not None:
        resolved_config.update({"max_jobs": int(max_jobs)})
    if queues:
        resolved_config.update({"queues": list(queues)})
    if verbosity:
        resolved_config.update({"verbosity": int(verbosity)})
    # queue -> queues
    if resolved_config.get("queue"):
        if isinstance(resolved_config.get("queue"), str):
            resolved_config["queues"].append(resolved_config["queue"])
        else:
            msg = (
                "Invalid launch agent config for key 'queue' with type: {type(resolved_config.get('queue'))} "
                "(expected str). Specify multiple queues with the 'queues' key"
            )
            raise LaunchError(msg)

    keys = ["entity"]
    settings = {
        k: resolved_config.get(k) for k in keys if resolved_config.get(k) is not None
    }

    api = Api(default_settings=settings)

    if resolved_config.get("entity") is None:
        resolved_config.update({"entity": api.default_entity})

    return resolved_config, api


def create_and_run_agent(
    api: Api,
    config: Dict[str, Any],
) -> None:
    try:
        from wandb.sdk.launch.agent import config as agent_config
    except ModuleNotFoundError:
        raise LaunchError(
            "wandb launch-agent requires pydantic to be installed. "
            "Please install with `pip install wandb[launch]`"
        )
    try:
        agent_config.AgentConfig(**config)
    except agent_config.ValidationError as e:
        errors = e.errors()
        for error in errors:
            loc = ".".join([str(x) for x in error.get("loc", [])])
            msg = f"Agent config error in field {loc}"
            value = error.get("input")
            if not isinstance(value, dict):
                msg += f" (value: {value})"
            msg += f": {error['msg']}"
            wandb.termerror(msg)
        raise LaunchError("Invalid launch agent config")
    agent = LaunchAgent(api, config)
    try:
        asyncio.run(agent.loop())
    except asyncio.CancelledError:
        pass


async def _launch(
    api: Api,
    job: Optional[str] = None,
    name: Optional[str] = None,
    project: Optional[str] = None,
    entity: Optional[str] = None,
    docker_image: Optional[str] = None,
    entry_point: Optional[List[str]] = None,
    version: Optional[str] = None,
    resource: Optional[str] = None,
    resource_args: Optional[Dict[str, Any]] = None,
    launch_config: Optional[Dict[str, Any]] = None,
    synchronous: Optional[bool] = None,
    run_id: Optional[str] = None,
    repository: Optional[str] = None,
) -> AbstractRun:
    """Helper that delegates to the project-running method corresponding to the passed-in backend."""
    if launch_config is None:
        launch_config = {}
    if resource is None:
        resource = "local-container"
    launch_spec = construct_launch_spec(
        None,
        job,
        api,
        name,
        project,
        entity,
        docker_image,
        resource,
        entry_point,
        version,
        resource_args,
        launch_config,
        run_id,
        repository,
        author=None,
    )
    validate_launch_spec_source(launch_spec)
    launch_project = LaunchProject.from_spec(launch_spec, api)
    launch_project.fetch_and_validate_project()
    entrypoint = launch_project.get_job_entry_point()
    image_uri = (
        launch_project.docker_image or launch_project.job_base_image
    )  # Either set by user or None.

    # construct runner config.
    runner_config: Dict[str, Any] = {}
    runner_config[PROJECT_SYNCHRONOUS] = synchronous

    config = launch_config or {}
    environment_config, build_config, registry_config = construct_agent_configs(config)
    environment = loader.environment_from_config(environment_config)
    if environment is not None and not isinstance(environment, LocalEnvironment):
        await environment.verify()
    registry = loader.registry_from_config(registry_config, environment)
    builder = loader.builder_from_config(build_config, environment, registry)
    if not (launch_project.docker_image or launch_project.job_base_image):
        assert entrypoint
        image_uri = await builder.build_image(launch_project, entrypoint, None)
    backend = loader.runner_from_config(
        resource, api, runner_config, environment, registry
    )
    if backend:
        assert image_uri
        submitted_run = await backend.run(launch_project, image_uri)
        # this check will always pass, run is only optional in the agent case where
        # a run queue id is present on the backend config
        assert submitted_run
        return submitted_run
    else:
        raise ExecutionError(
            f"Unavailable backend {resource}, available backends: {', '.join(loader.WANDB_RUNNERS)}"
        )


def launch(
    api: Api,
    job: Optional[str] = None,
    entry_point: Optional[List[str]] = None,
    version: Optional[str] = None,
    name: Optional[str] = None,
    resource: Optional[str] = None,
    resource_args: Optional[Dict[str, Any]] = None,
    project: Optional[str] = None,
    entity: Optional[str] = None,
    docker_image: Optional[str] = None,
    config: Optional[Dict[str, Any]] = None,
    synchronous: Optional[bool] = True,
    run_id: Optional[str] = None,
    repository: Optional[str] = None,
) -> AbstractRun:
    """Launch a W&B launch experiment.

    Arguments:
        job: string reference to a wandb.Job eg: wandb/test/my-job:latest
        api: An instance of a wandb Api from wandb.apis.internal.
        entry_point: Entry point to run within the project. Defaults to using the entry point used
            in the original run for wandb URIs, or main.py for git repository URIs.
        version: For Git-based projects, either a commit hash or a branch name.
        name: Name run under which to launch the run.
        resource: Execution backend for the run.
        resource_args: Resource related arguments for launching runs onto a remote backend.
            Will be stored on the constructed launch config under ``resource_args``.
        project: Target project to send launched run to
        entity: Target entity to send launched run to
        config: A dictionary containing the configuration for the run. May also contain
        resource specific arguments under the key "resource_args".
        synchronous: Whether to block while waiting for a run to complete. Defaults to True.
            Note that if ``synchronous`` is False and ``backend`` is "local-container", this
            method will return, but the current process will block when exiting until
            the local run completes. If the current process is interrupted, any
            asynchronous runs launched via this method will be terminated. If
            ``synchronous`` is True and the run fails, the current process will
            error out as well.
        run_id: ID for the run (To ultimately replace the :name: field)
        repository: string name of repository path for remote registry

    Example:
        ```python
        from wandb.sdk.launch import launch

        job = "wandb/jobs/Hello World:latest"
        params = {"epochs": 5}
        # Run W&B project and create a reproducible docker environment
        # on a local host
        api = wandb.apis.internal.Api()
        launch(api, job, parameters=params)
        ```


    Returns:
        an instance of`wandb.launch.SubmittedRun` exposing information (e.g. run ID)
        about the launched run.

    Raises:
        `wandb.exceptions.ExecutionError` If a run launched in blocking mode
        is unsuccessful.
    """
    submitted_run_obj = asyncio.run(
        _launch(
            job=job,
            name=name,
            project=project,
            entity=entity,
            docker_image=docker_image,
            entry_point=entry_point,
            version=version,
            resource=resource,
            resource_args=resource_args,
            launch_config=config,
            synchronous=synchronous,
            api=api,
            run_id=run_id,
            repository=repository,
        )
    )

    return submitted_run_obj