wandb 0.16.3__py3-none-any.whl → 0.16.5__py3-none-any.whl

wandb/__init__.py

@@ -11,8 +11,8 @@ For scripts and interactive notebooks, see https://github.com/wandb/examples.
 
 For reference documentation, see https://docs.wandb.com/ref/python.
 """
-__version__ = "0.16.3"
-_minimum_core_version = "0.17.0b8"
+__version__ = "0.16.5"
+_minimum_core_version = "0.17.0b10"
 
 # Used with pypi checks and other messages related to pip
 _wandb_module = "wandb"

wandb/agents/pyagent.py

@@ -347,7 +347,7 @@ def pyagent(sweep_id, function, entity=None, project=None, count=None):
         count (int, optional): the number of trials to run.
     """
     if not callable(function):
-        raise Exception("function paramter must be callable!")
+        raise Exception("function parameter must be callable!")
     agent = Agent(
         sweep_id,
         function=function,

wandb/apis/importers/__init__.py

@@ -1,4 +1 @@
-from wandb.util import get_module
-
-if get_module("mlflow"):
-    from .mlflow import MlflowImporter, MlflowRun  # noqa: F401
+from .internals.util import Namespace

wandb/apis/importers/internals/internal.py

@@ -0,0 +1,386 @@
+import json
+import logging
+import math
+import os
+import queue
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional
+
+import numpy as np
+from google.protobuf.json_format import ParseDict
+from tenacity import retry, stop_after_attempt, wait_random_exponential
+
+from wandb import Artifact
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_settings_pb2
+from wandb.proto import wandb_telemetry_pb2 as telem_pb
+from wandb.sdk.interface.interface import file_policy_to_enum
+from wandb.sdk.interface.interface_queue import InterfaceQueue
+from wandb.sdk.internal import context
+from wandb.sdk.internal.sender import SendManager
+from wandb.sdk.internal.settings_static import SettingsStatic
+from wandb.util import coalesce, recursive_cast_dictlike_to_dict
+
+from .protocols import ImporterRun
+
+ROOT_DIR = "./wandb-importer"
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+if os.getenv("WANDB_IMPORTER_ENABLE_RICH_LOGGING"):
+    from rich.logging import RichHandler
+
+    logger.addHandler(RichHandler(rich_tracebacks=True, tracebacks_show_locals=True))
+else:
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.INFO)
+
+    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+    console_handler.setFormatter(formatter)
+
+    logger.addHandler(console_handler)
+
+
+exp_retry = retry(
+    wait=wait_random_exponential(multiplier=1, max=10), stop=stop_after_attempt(3)
+)
+
+
+class AlternateSendManager(SendManager):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._send_artifact = exp_retry(self._send_artifact)
+
+
+@dataclass(frozen=True)
+class SendManagerConfig:
+    """Configure which parts of SendManager tooling to use."""
+
+    use_artifacts: bool = False
+    log_artifacts: bool = False
+    metadata: bool = False
+    files: bool = False
+    media: bool = False
+    code: bool = False
+    history: bool = False
+    summary: bool = False
+    terminal_output: bool = False
+
+
+@dataclass
+class RecordMaker:
+    run: ImporterRun
+    interface: InterfaceQueue = InterfaceQueue()
+
+    @property
+    def run_dir(self) -> str:
+        p = Path(f"{ROOT_DIR}/{self.run.run_id()}/wandb")
+        p.mkdir(parents=True, exist_ok=True)
+        return f"{ROOT_DIR}/{self.run.run_id()}"
+
+    def make_artifacts_only_records(
+        self,
+        artifacts: Optional[Iterable[Artifact]] = None,
+        used_artifacts: Optional[Iterable[Artifact]] = None,
+    ) -> Iterable[pb.Record]:
+        """Only make records required to upload artifacts.
+
+        Escape hatch for adding extra artifacts to a run.
+        """
+        yield self._make_run_record()
+
+        if used_artifacts:
+            for art in used_artifacts:
+                yield self._make_artifact_record(art, use_artifact=True)
+
+        if artifacts:
+            for art in artifacts:
+                yield self._make_artifact_record(art)
+
+    def make_records(
+        self,
+        config: SendManagerConfig,
+    ) -> Iterable[pb.Record]:
+        """Make all the records that constitute a run."""
+        yield self._make_run_record()
+        yield self._make_telem_record()
+
+        include_artifacts = config.log_artifacts or config.use_artifacts
+        yield self._make_files_record(
+            include_artifacts, config.files, config.media, config.code
+        )
+
+        if config.use_artifacts:
+            if (used_artifacts := self.run.used_artifacts()) is not None:
+                for artifact in used_artifacts:
+                    yield self._make_artifact_record(artifact, use_artifact=True)
+
+        if config.log_artifacts:
+            if (artifacts := self.run.artifacts()) is not None:
+                for artifact in artifacts:
+                    yield self._make_artifact_record(artifact)
+
+        if config.history:
+            yield from self._make_history_records()
+
+        if config.summary:
+            yield self._make_summary_record()
+
+        if config.terminal_output:
+            if (lines := self.run.logs()) is not None:
+                for line in lines:
+                    yield self._make_output_record(line)
+
+    def _make_run_record(self) -> pb.Record:
+        run = pb.RunRecord()
+        run.run_id = self.run.run_id()
+        run.entity = self.run.entity()
+        run.project = self.run.project()
+        run.display_name = coalesce(self.run.display_name())
+        run.notes = coalesce(self.run.notes(), "")
+        run.tags.extend(coalesce(self.run.tags(), []))
+        run.start_time.FromMilliseconds(self.run.start_time())
+
+        host = self.run.host()
+        if host is not None:
+            run.host = host
+
+        runtime = self.run.runtime()
+        if runtime is not None:
+            run.runtime = runtime
+
+        run_group = self.run.run_group()
+        if run_group is not None:
+            run.run_group = run_group
+
+        config = self.run.config()
+        if "_wandb" not in config:
+            config["_wandb"] = {}
+
+        # how do I get this automatically?
+        config["_wandb"]["code_path"] = self.run.code_path()
+        config["_wandb"]["python_version"] = self.run.python_version()
+        config["_wandb"]["cli_version"] = self.run.cli_version()
+
+        self.interface._make_config(
+            data=config,
+            obj=run.config,
+        )  # is there a better way?
+        return self.interface._make_record(run=run)
+
+    def _make_output_record(self, line) -> pb.Record:
+        output_raw = pb.OutputRawRecord()
+        output_raw.output_type = pb.OutputRawRecord.OutputType.STDOUT
+        output_raw.line = line
+        return self.interface._make_record(output_raw=output_raw)
+
+    def _make_summary_record(self) -> pb.Record:
+        d: dict = {
+            **self.run.summary(),
+            "_runtime": self.run.runtime(),  # quirk of runtime -- it has to be here!
+            # '_timestamp': self.run.start_time()/1000,
+        }
+        d = recursive_cast_dictlike_to_dict(d)
+        summary = self.interface._make_summary_from_dict(d)
+        return self.interface._make_record(summary=summary)
+
+    def _make_history_records(self) -> Iterable[pb.Record]:
+        for metrics in self.run.metrics():
+            history = pb.HistoryRecord()
+            for k, v in metrics.items():
+                item = history.item.add()
+                item.key = k
+                # There seems to be some conversion issue to breaks when we try to re-upload.
+                # np.NaN gets converted to float("nan"), which is not expected by our system.
+                # If this cast to string (!) is not done, the row will be dropped.
+                if (isinstance(v, float) and math.isnan(v)) or v == "NaN":
+                    v = np.NaN
+
+                if isinstance(v, bytes):
+                    # it's a json string encoded as bytes
+                    v = v.decode("utf-8")
+                else:
+                    v = json.dumps(v)
+
+                item.value_json = v
+            rec = self.interface._make_record(history=history)
+            yield rec
+
+    def _make_files_record(
+        self, artifacts: bool, files: bool, media: bool, code: bool
+    ) -> pb.Record:
+        run_files = self.run.files()
+        metadata_fname = f"{self.run_dir}/files/wandb-metadata.json"
+        if not files or run_files is None:
+            # We'll always need a metadata file even if there are no other files to upload
+            metadata_fname = self._make_metadata_file()
+            run_files = [(metadata_fname, "end")]
+        files_record = pb.FilesRecord()
+        for path, policy in run_files:
+            if not artifacts and path.startswith("artifact/"):
+                continue
+            if not media and path.startswith("media/"):
+                continue
+            if not code and path.startswith("code/"):
+                continue
+
+            # DirWatcher requires the path to start with media/ instead of the full path
+            if "media" in path:
+                p = Path(path)
+                path = str(p.relative_to(f"{self.run_dir}/files"))
+            f = files_record.files.add()
+            f.path = path
+            f.policy = file_policy_to_enum(policy)
+
+        return self.interface._make_record(files=files_record)
+
+    def _make_artifact_record(
+        self, artifact: Artifact, use_artifact=False
+    ) -> pb.Record:
+        proto = self.interface._make_artifact(artifact)
+        proto.run_id = str(self.run.run_id())
+        proto.project = str(self.run.project())
+        proto.entity = str(self.run.entity())
+        proto.user_created = use_artifact
+        proto.use_after_commit = use_artifact
+        proto.finalize = True
+
+        aliases = artifact._aliases
+        aliases += ["latest", "imported"]
+
+        for alias in aliases:
+            proto.aliases.append(alias)
+        return self.interface._make_record(artifact=proto)
+
+    def _make_telem_record(self) -> pb.Record:
+        telem = telem_pb.TelemetryRecord()
+
+        feature = telem_pb.Feature()
+        feature.importer_mlflow = True
+        telem.feature.CopyFrom(feature)
+
+        cli_version = self.run.cli_version()
+        if cli_version:
+            telem.cli_version = cli_version
+
+        python_version = self.run.python_version()
+        if python_version:
+            telem.python_version = python_version
+
+        return self.interface._make_record(telemetry=telem)
+
+    def _make_metadata_file(self) -> str:
+        missing_text = "This data was not captured"
+        files_dir = f"{self.run_dir}/files"
+        os.makedirs(files_dir, exist_ok=True)
+
+        d = {}
+        d["os"] = coalesce(self.run.os_version(), missing_text)
+        d["python"] = coalesce(self.run.python_version(), missing_text)
+        d["program"] = coalesce(self.run.program(), missing_text)
+        d["cuda"] = coalesce(self.run.cuda_version(), missing_text)
+        d["host"] = coalesce(self.run.host(), missing_text)
+        d["username"] = coalesce(self.run.username(), missing_text)
+        d["executable"] = coalesce(self.run.executable(), missing_text)
+
+        gpus_used = self.run.gpus_used()
+        if gpus_used is not None:
+            d["gpu_devices"] = json.dumps(gpus_used)
+            d["gpu_count"] = json.dumps(len(gpus_used))
+
+        cpus_used = self.run.cpus_used()
+        if cpus_used is not None:
+            d["cpu_count"] = json.dumps(self.run.cpus_used())
+
+        mem_used = self.run.memory_used()
+        if mem_used is not None:
+            d["memory"] = json.dumps({"total": self.run.memory_used()})
+
+        fname = f"{files_dir}/wandb-metadata.json"
+        with open(fname, "w") as f:
+            f.write(json.dumps(d))
+        return fname
+
+
+def _make_settings(
+    root_dir: str, settings_override: Optional[Dict[str, Any]] = None
+) -> SettingsStatic:
+    _settings_override = coalesce(settings_override, {})
+
+    default_settings: Dict[str, Any] = {
+        "files_dir": os.path.join(root_dir, "files"),
+        "root_dir": root_dir,
+        "sync_file": os.path.join(root_dir, "txlog.wandb"),
+        "resume": "false",
+        "program": None,
+        "ignore_globs": [],
+        "disable_job_creation": True,
+        "_start_time": 0,
+        "_offline": None,
+        "_sync": True,
+        "_live_policy_rate_limit": 15,  # matches dir_watcher
+        "_live_policy_wait_time": 600,  # matches dir_watcher
+        "_async_upload_concurrency_limit": None,
+        "_file_stream_timeout_seconds": 60,
+    }
+
+    combined_settings = {**default_settings, **_settings_override}
+    settings_message = wandb_settings_pb2.Settings()
+    ParseDict(combined_settings, settings_message)
+
+    return SettingsStatic(settings_message)
+
+
+def send_run(
+    run: ImporterRun,
+    *,
+    extra_arts: Optional[Iterable[Artifact]] = None,
+    extra_used_arts: Optional[Iterable[Artifact]] = None,
+    config: Optional[SendManagerConfig] = None,
+    overrides: Optional[Dict[str, Any]] = None,
+    settings_override: Optional[Dict[str, Any]] = None,
+) -> None:
+    if config is None:
+        config = SendManagerConfig()
+
+    # does this need to be here for pmap?
+    if overrides:
+        for k, v in overrides.items():
+            # `lambda: v` won't work!
+            # https://stackoverflow.com/questions/10802002/why-deepcopy-doesnt-create-new-references-to-lambda-function
+            setattr(run, k, lambda v=v: v)
+
+    rm = RecordMaker(run)
+    root_dir = rm.run_dir
+
+    settings = _make_settings(root_dir, settings_override)
+    sm_record_q = queue.Queue()
+    # wm_record_q = queue.Queue()
+    result_q = queue.Queue()
+    interface = InterfaceQueue(record_q=sm_record_q)
+    context_keeper = context.ContextKeeper()
+    sm = AlternateSendManager(
+        settings, sm_record_q, result_q, interface, context_keeper
+    )
+    # wm = WriteManager(
+    #     settings, wm_record_q, result_q, sm_record_q, interface, context_keeper
+    # )
+
+    if extra_arts or extra_used_arts:
+        records = rm.make_artifacts_only_records(extra_arts, extra_used_arts)
+    else:
+        records = rm.make_records(config)
+
+    for r in records:
+        logger.debug(f"Sending {r=}")
+        # In a future update, it might be good to write to a transaction log and have
+        # incremental uploads only send the missing records.
+        # wm.write(r)
+
+        sm.send(r)
+
+    sm.finish()
+    # wm.finish()

wandb/apis/importers/internals/protocols.py

@@ -0,0 +1,125 @@
+import logging
+import sys
+from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple
+
+from wandb.sdk.artifacts.artifact import Artifact
+
+if sys.version_info >= (3, 8):
+    from typing import Protocol, runtime_checkable
+else:
+    from typing_extensions import Protocol, runtime_checkable
+
+logger = logging.getLogger("import_logger")
+
+PathStr = str
+Policy = Literal["now", "end", "live"]
+
+
+@runtime_checkable
+class ImporterRun(Protocol):
+    def run_id(self) -> str:
+        ...  # pragma: no cover
+
+    def entity(self) -> str:
+        ...  # pragma: no cover
+
+    def project(self) -> str:
+        ...  # pragma: no cover
+
+    def config(self) -> Dict[str, Any]:
+        ...  # pragma: no cover
+
+    def summary(self) -> Dict[str, float]:
+        ...  # pragma: no cover
+
+    def metrics(self) -> Iterable[Dict[str, float]]:
+        """Metrics for the run.
+
+        We expect metrics in this shape:
+
+        [
+            {'metric1': 1, 'metric2': 1, '_step': 0},
+            {'metric1': 2, 'metric2': 4, '_step': 1},
+            {'metric1': 3, 'metric2': 9, '_step': 2},
+            ...
+        ]
+
+        You can also submit metrics in this shape:
+        [
+            {'metric1': 1, '_step': 0},
+            {'metric2': 1, '_step': 0},
+            {'metric1': 2, '_step': 1},
+            {'metric2': 4, '_step': 1},
+            ...
+        ]
+        """
+        ...  # pragma: no cover
+
+    def run_group(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def job_type(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def display_name(self) -> str:
+        ...  # pragma: no cover
+
+    def notes(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def tags(self) -> Optional[List[str]]:
+        ...  # pragma: no cover
+
+    def artifacts(self) -> Optional[Iterable[Artifact]]:
+        ...  # pragma: no cover
+
+    def used_artifacts(self) -> Optional[Iterable[Artifact]]:
+        ...  # pragma: no cover
+
+    def os_version(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def python_version(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def cuda_version(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def program(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def host(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def username(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def executable(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def gpus_used(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def cpus_used(self) -> Optional[int]:
+        ...  # pragma: no cover
+
+    def memory_used(self) -> Optional[int]:
+        ...  # pragma: no cover
+
+    def runtime(self) -> Optional[int]:
+        ...  # pragma: no cover
+
+    def start_time(self) -> Optional[int]:
+        ...  # pragma: no cover
+
+    def code_path(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def cli_version(self) -> Optional[str]:
+        ...  # pragma: no cover
+
+    def files(self) -> Optional[Iterable[Tuple[PathStr, Policy]]]:
+        ...  # pragma: no cover
+
+    def logs(self) -> Optional[Iterable[str]]:
+        ...  # pragma: no cover

wandb/apis/importers/internals/util.py

@@ -0,0 +1,78 @@
+import logging
+import sys
+import traceback
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+
+@dataclass(frozen=True)
+class Namespace:
+    """Configure an alternate entity/project at the dst server your data will end up in."""
+
+    entity: str
+    project: str
+
+    @classmethod
+    def from_path(cls, path: str):
+        entity, project = path.split("/")
+        return cls(entity, project)
+
+    @property
+    def path(self):
+        return f"{self.entity}/{self.project}"
+
+    @property
+    def send_manager_overrides(self):
+        overrides = {}
+        if self.entity:
+            overrides["entity"] = self.entity
+        if self.project:
+            overrides["project"] = self.project
+        return overrides
+
+
+logger = logging.getLogger("import_logger")
+
+
+def parallelize(
+    func,
+    iterable: Iterable,
+    *args,
+    max_workers: Optional[int] = None,
+    raise_on_error: bool = False,
+    **kwargs,
+):
+    def safe_func(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            _, _, exc_traceback = sys.exc_info()
+            traceback_details = traceback.extract_tb(exc_traceback)
+            filename = traceback_details[-1].filename
+            lineno = traceback_details[-1].lineno
+            logger.debug(
+                f"Exception: {func=} {args=} {kwargs=} {e=} {filename=} {lineno=}. {traceback_details=}"
+            )
+            if raise_on_error:
+                raise e
+
+    results = []
+    with ThreadPoolExecutor(max_workers) as exc:
+        futures = {exc.submit(safe_func, x, *args, **kwargs): x for x in iterable}
+        for future in as_completed(futures):
+            results.append(future.result())
+    return results
+
+
+def for_each(
+    func, iterable: Iterable, parallel: bool = True, max_workers: Optional[int] = None
+):
+    if parallel:
+        return parallelize(
+            func,
+            iterable,
+            max_workers=max_workers,
+        )
+
+    return [func(x) for x in iterable]