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

wandb/apis/importers/base.py

@@ -1,400 +0,0 @@
-import json
-import os
-import queue
-import sys
-from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List, Optional, Tuple
-from unittest.mock import patch
-
-from google.protobuf.json_format import ParseDict
-from tqdm import tqdm
-
-import wandb
-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 cast_dictlike_to_dict, coalesce
-
-if sys.version_info >= (3, 8):
-    from typing import Protocol
-else:
-    from typing_extensions import Protocol
-
-
-with patch("click.echo"):
-    from wandb.apis.reports import Report
-
-
-class ImporterRun(Protocol):
-    def run_id(self) -> str:
-        ...
-
-    def entity(self) -> str:
-        ...
-
-    def project(self) -> str:
-        ...
-
-    def config(self) -> Dict[str, Any]:
-        ...
-
-    def summary(self) -> Dict[str, float]:
-        ...
-
-    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},
-            ...
-        ]
-        """
-        ...
-
-    def run_group(self) -> Optional[str]:
-        ...
-
-    def job_type(self) -> Optional[str]:
-        ...
-
-    def display_name(self) -> str:
-        ...
-
-    def notes(self) -> Optional[str]:
-        ...
-
-    def tags(self) -> Optional[List[str]]:
-        ...
-
-    def artifacts(self) -> Optional[Iterable[wandb.Artifact]]:  # type: ignore
-        ...
-
-    def used_artifacts(self) -> Optional[Iterable[wandb.Artifact]]:  # type: ignore
-        ...
-
-    def os_version(self) -> Optional[str]:
-        ...
-
-    def python_version(self) -> Optional[str]:
-        ...
-
-    def cuda_version(self) -> Optional[str]:
-        ...
-
-    def program(self) -> Optional[str]:
-        ...
-
-    def host(self) -> Optional[str]:
-        ...
-
-    def username(self) -> Optional[str]:
-        ...
-
-    def executable(self) -> Optional[str]:
-        ...
-
-    def gpus_used(self) -> Optional[str]:
-        ...
-
-    def cpus_used(self) -> Optional[int]:  # can we get the model?
-        ...
-
-    def memory_used(self) -> Optional[int]:
-        ...
-
-    def runtime(self) -> Optional[int]:
-        ...
-
-    def start_time(self) -> Optional[int]:
-        ...
-
-    def code_path(self) -> Optional[str]:
-        ...
-
-    def cli_version(self) -> Optional[str]:
-        ...
-
-    def files(self) -> Optional[Iterable[Tuple[str, str]]]:
-        ...
-
-    def logs(self) -> Optional[Iterable[str]]:
-        ...
-
-
-class Importer(Protocol):
-    def collect_runs(self, *args, **kwargs) -> Iterable[ImporterRun]:
-        ...
-
-    def collect_reports(self, *args, **kwargs) -> Iterable[Report]:
-        ...
-
-    def import_run(self, run: ImporterRun) -> None:
-        ...
-
-    def import_report(self, report: Report) -> None:
-        ...
-
-
-@dataclass
-class RecordMaker:
-    run: ImporterRun
-    interface: InterfaceQueue = InterfaceQueue()
-
-    @property
-    def run_dir(self) -> str:
-        return f"./wandb-importer/{self.run.run_id()}"
-
-    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 = 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 enumerate(self.run.metrics()):
-            history = pb.HistoryRecord()
-            for k, v in metrics.items():
-                item = history.item.add()
-                item.key = k
-                item.value_json = json.dumps(v)
-            yield self.interface._make_record(history=history)
-
-    def _make_files_record(
-        self,
-        files_dict,
-    ) -> pb.Record:
-        files_record = pb.FilesRecord()
-        for path, policy in files_dict["files"]:
-            f = files_record.files.add()
-            f.path = path
-            f.policy = file_policy_to_enum(policy)  # is this always "end"?
-        return self.interface._make_record(files=files_record)
-
-    def _make_metadata_files_record(self) -> pb.Record:
-        files = self.run.files()
-        if files is None:
-            metadata_fname = self._make_metadata_file()
-            files = [(metadata_fname, "end")]
-
-        files_dict = {"files": files}
-        return self._make_files_record(files_dict)
-
-    def _make_artifact_record(self, artifact, use_artifact=False) -> pb.Record:
-        proto = self.interface._make_artifact(artifact)
-        proto.run_id = self.run.run_id()
-        proto.project = self.run.project()
-        proto.entity = self.run.entity()
-        proto.user_created = use_artifact
-        proto.use_after_commit = use_artifact
-        proto.finalize = True
-        for tag in ["latest", "imported"]:
-            proto.aliases.append(tag)
-        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"
-
-        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"{self.run_dir}/files/wandb-metadata.json"
-        with open(fname, "w") as f:
-            f.write(json.dumps(d))
-        return fname
-
-
-def send_run_with_send_manager(
-    run: ImporterRun,
-    overrides: Optional[Dict[str, Any]] = None,
-    settings_override: Optional[SettingsStatic] = None,
-) -> None:
-    # 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)
-    _settings_override = coalesce(settings_override, {})
-    rm = RecordMaker(run)
-
-    root_dir = rm.run_dir
-    default_settings = {
-        "files_dir": os.path.join(root_dir, "files"),
-        "root_dir": root_dir,
-        "_start_time": 0,
-        "git_remote": None,
-        "resume": False,
-        "program": None,
-        "ignore_globs": (),
-        "run_id": None,
-        "entity": None,
-        "project": None,
-        "run_group": None,
-        "job_type": None,
-        "run_tags": None,
-        "run_name": None,
-        "run_notes": None,
-        "save_code": None,
-        "email": None,
-        "silent": None,
-        "_offline": None,
-        "_sync": True,
-        "_live_policy_rate_limit": None,
-        "_live_policy_wait_time": None,
-        "disable_job_creation": False,
-        "_async_upload_concurrency_limit": None,
-    }
-    combined_settings = {**default_settings, **_settings_override}
-    settings_message = wandb_settings_pb2.Settings()
-    ParseDict(combined_settings, settings_message)
-
-    settings = SettingsStatic(settings_message)
-
-    record_q: queue.Queue = queue.Queue()
-    result_q: queue.Queue = queue.Queue()
-    interface = InterfaceQueue(record_q=record_q)
-    context_keeper = context.ContextKeeper()
-
-    with SendManager(settings, record_q, result_q, interface, context_keeper) as sm:
-        wandb.termlog(">> Make run record")
-        sm.send(rm._make_run_record())
-
-        wandb.termlog(">> Use Artifacts")
-        used_artifacts = run.used_artifacts()
-        if used_artifacts is not None:
-            for artifact in tqdm(
-                used_artifacts, desc="Used artifacts", unit="artifacts", leave=False
-            ):
-                sm.send(rm._make_artifact_record(artifact, use_artifact=True))
-
-        wandb.termlog(">> Log Artifacts")
-        artifacts = run.artifacts()
-        if artifacts is not None:
-            for artifact in tqdm(
-                artifacts, desc="Logged artifacts", unit="artifacts", leave=False
-            ):
-                sm.send(rm._make_artifact_record(artifact))
-
-        wandb.termlog(">> Log Metadata")
-        sm.send(rm._make_metadata_files_record())
-
-        wandb.termlog(">> Log History")
-        for history_record in tqdm(
-            rm._make_history_records(), desc="History", unit="steps", leave=False
-        ):
-            sm.send(history_record)
-
-        wandb.termlog(">> Log Summary")
-        sm.send(rm._make_summary_record())
-
-        wandb.termlog(">> Log Output")
-        # if hasattr(run, "_logs"):
-        #     lines = run._logs
-        lines = run.logs()
-        if lines is not None:
-            for line in tqdm(lines, desc="Stdout", unit="lines", leave=False):
-                sm.send(rm._make_output_record(line))
-
-        wandb.termlog(">> Log Telem")
-        sm.send(rm._make_telem_record())