wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/internal/writer.py

@@ -0,0 +1,204 @@
+"""Writer thread."""
+
+import logging
+from typing import TYPE_CHECKING, Callable, Optional
+
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_telemetry_pb2 as tpb
+
+from ..interface.interface_queue import InterfaceQueue
+from ..lib import proto_util, telemetry
+from . import context, datastore, flow_control
+from .settings_static import SettingsStatic
+
+if TYPE_CHECKING:
+    from queue import Queue
+
+
+logger = logging.getLogger(__name__)
+
+
+class WriteManager:
+    _settings: SettingsStatic
+    _record_q: "Queue[pb.Record]"
+    _result_q: "Queue[pb.Result]"
+    _sender_q: "Queue[pb.Record]"
+    _interface: InterfaceQueue
+    _context_keeper: context.ContextKeeper
+
+    _ds: Optional[datastore.DataStore]
+    _flow_control: Optional[flow_control.FlowControl]
+    _status_report: Optional["pb.StatusReportRequest"]
+    _record_num: int
+    _telemetry_obj: tpb.TelemetryRecord
+    _telemetry_overflow: bool
+    _use_flow_control: bool
+
+    # TODO(cancel_paused): implement me
+    # _sender_cancel_set: Set[str]
+
+    def __init__(
+        self,
+        settings: SettingsStatic,
+        record_q: "Queue[pb.Record]",
+        result_q: "Queue[pb.Result]",
+        sender_q: "Queue[pb.Record]",
+        interface: InterfaceQueue,
+        context_keeper: context.ContextKeeper,
+    ):
+        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
+
+        # TODO(cancel_paused): implement me
+        # self._sender_cancel_set = set()
+
+        self._ds = None
+        self._flow_control = None
+        self._status_report = None
+        self._record_num = 0
+        self._telemetry_obj = tpb.TelemetryRecord()
+        self._telemetry_overflow = False
+        self._use_flow_control = not (
+            self._settings.x_flow_control_disabled or self._settings._offline
+        )
+
+    def open(self) -> None:
+        self._ds = datastore.DataStore()
+        self._ds.open_for_write(self._settings.sync_file)
+        self._flow_control = flow_control.FlowControl(
+            settings=self._settings,
+            write_record=self._write_record,
+            forward_record=self._forward_record,
+            pause_marker=self._pause_marker,
+            recover_records=self._recover_records,
+        )
+
+    def _forward_record(self, record: "pb.Record") -> None:
+        self._context_keeper.add_from_record(record)
+        self._sender_q.put(record)
+
+    def _send_mark(self) -> None:
+        sender_mark = pb.SenderMarkRequest()
+        record = self._interface._make_request(sender_mark=sender_mark)
+        self._forward_record(record)
+
+    def _maybe_send_telemetry(self) -> None:
+        if self._telemetry_overflow:
+            return
+        self._telemetry_overflow = True
+        with telemetry.context(obj=self._telemetry_obj) as tel:
+            tel.feature.flow_control_overflow = True
+        telemetry_record = pb.TelemetryRecordRequest(telemetry=self._telemetry_obj)
+        record = self._interface._make_request(telemetry_record=telemetry_record)
+        self._forward_record(record)
+
+    def _pause_marker(self) -> None:
+        self._maybe_send_telemetry()
+        self._send_mark()
+
+    def _write_record(self, record: "pb.Record") -> int:
+        assert self._ds
+
+        self._record_num += 1
+        proto_util._assign_record_num(record, self._record_num)
+        ret = self._ds.write(record)
+        assert ret is not None
+
+        _start_offset, end_offset, _flush_offset = ret
+        proto_util._assign_end_offset(record, end_offset)
+        return end_offset
+
+    def _ensure_flushed(self, offset: int) -> None:
+        if self._ds:
+            self._ds.ensure_flushed(offset)
+
+    def _recover_records(self, start: int, end: int) -> None:
+        sender_read = pb.SenderReadRequest(start_offset=start, final_offset=end)
+        # TODO(cancel_paused): implement me
+        # for cancel_id in self._sender_cancel_set:
+        #     sender_read.cancel_list.append(cancel_id)
+        record = self._interface._make_request(sender_read=sender_read)
+        self._ensure_flushed(end)
+        self._forward_record(record)
+
+    def _write(self, record: "pb.Record") -> None:
+        if not self._ds:
+            self.open()
+        assert self._flow_control
+
+        if not record.control.local:
+            self._write_record(record)
+
+        if self._use_flow_control:
+            self._flow_control.flow(record)
+        elif not self._settings._offline or record.control.always_send:
+            # when flow_control is disabled we pass through all records to
+            # the sender as long as we are online.  The exception is there
+            # are special records that we always pass to the sender
+            # (namely the exit record so we can trigger the defer shutdown
+            # state machine)
+            self._forward_record(record)
+
+    def write(self, record: "pb.Record") -> None:
+        record_type = record.WhichOneof("record_type")
+        assert record_type
+        writer_str = "write_" + record_type
+        write_handler: Callable[[pb.Record], None] = getattr(
+            self, writer_str, self._write
+        )
+        write_handler(record)
+
+    def write_request(self, record: "pb.Record") -> None:
+        request_type = record.request.WhichOneof("request_type")
+        assert request_type
+        write_request_str = "write_request_" + request_type
+        write_request_handler: Optional[Callable[[pb.Record], None]] = getattr(
+            self, write_request_str, None
+        )
+        if write_request_handler:
+            return write_request_handler(record)
+        self._write(record)
+
+    def write_request_run_status(self, record: "pb.Record") -> None:
+        result = proto_util._result_from_record(record)
+        if self._status_report:
+            result.response.run_status_response.sync_time.CopyFrom(
+                self._status_report.sync_time
+            )
+            send_record_num = self._status_report.record_num
+            result.response.run_status_response.sync_items_total = self._record_num
+            result.response.run_status_response.sync_items_pending = (
+                self._record_num - send_record_num
+            )
+        self._respond_result(result)
+
+    def write_request_status_report(self, record: "pb.Record") -> None:
+        self._status_report = record.request.status_report
+        self._write(record)
+
+    def write_request_cancel(self, record: "pb.Record") -> None:
+        cancel_id = record.request.cancel.cancel_slot
+        self._context_keeper.cancel(cancel_id)
+
+        # TODO(cancel_paused): implement me
+        # cancelled = self._context_keeper.cancel(cancel_id)
+        # if not cancelled:
+        #     self._sender_cancel_set.add(cancel_id)
+
+    def _respond_result(self, result: "pb.Result") -> None:
+        self._result_q.put(result)
+
+    def finish(self) -> None:
+        if self._flow_control:
+            self._flow_control.flush()
+        if self._ds:
+            self._ds.close()
+        # TODO(debug_context) see context.py
+        # self._context_keeper._debug_print_orphans(print_to_stdout=self._settings._debug)
+
+    def debounce(self) -> None:
+        pass

wandb/sdk/launch/__init__.py

@@ -0,0 +1,15 @@
+from ._launch import create_and_run_agent, launch
+from ._launch_add import launch_add
+from .agent.agent import LaunchAgent
+from .inputs.manage import manage_config_file, manage_wandb_config
+from .utils import load_wandb_config
+
+__all__ = [
+    "create_and_run_agent",
+    "LaunchAgent",
+    "launch",
+    "launch_add",
+    "load_wandb_config",
+    "manage_config_file",
+    "manage_wandb_config",
+]

wandb/sdk/launch/_launch.py

@@ -0,0 +1,331 @@
+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(  # noqa: C901
+    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