wandb 0.15.9__py3-none-any.whl → 0.15.11__py3-none-any.whl

wandb/proto/wandb_server_pb2_grpc.py

@@ -1,8 +0,0 @@
-import google.protobuf
-
-protobuf_version = google.protobuf.__version__[0]
-
-if protobuf_version == "3":
-    from wandb.proto.v3.wandb_server_pb2_grpc import *
-elif protobuf_version == "4":
-    from wandb.proto.v4.wandb_server_pb2_grpc import *

wandb/sdk/artifacts/storage_policies/s3_bucket_policy.py

@@ -1,61 +0,0 @@
-"""S3 bucket storage policy."""
-from typing import TYPE_CHECKING, Dict, Optional, Sequence, Union
-
-from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
-from wandb.sdk.artifacts.storage_handlers.local_file_handler import LocalFileHandler
-from wandb.sdk.artifacts.storage_handlers.multi_handler import MultiHandler
-from wandb.sdk.artifacts.storage_handlers.s3_handler import S3Handler
-from wandb.sdk.artifacts.storage_handlers.tracking_handler import TrackingHandler
-from wandb.sdk.artifacts.storage_policy import StoragePolicy
-
-if TYPE_CHECKING:
-    from wandb.sdk.artifacts.artifact import Artifact
-    from wandb.sdk.lib.paths import FilePathStr, URIStr
-
-
-# Don't use this yet!
-class __S3BucketPolicy(StoragePolicy):  # noqa: N801
-    @classmethod
-    def name(cls) -> str:
-        return "wandb-s3-bucket-policy-v1"
-
-    @classmethod
-    def from_config(cls, config: Dict[str, str]) -> "__S3BucketPolicy":
-        if "bucket" not in config:
-            raise ValueError("Bucket name not found in config")
-        return cls(config["bucket"])
-
-    def __init__(self, bucket: str) -> None:
-        self._bucket = bucket
-        s3 = S3Handler(bucket)
-        local = LocalFileHandler()
-
-        self._handler = MultiHandler(
-            handlers=[
-                s3,
-                local,
-            ],
-            default_handler=TrackingHandler(),
-        )
-
-    def config(self) -> Dict[str, str]:
-        return {"bucket": self._bucket}
-
-    def load_path(
-        self,
-        manifest_entry: ArtifactManifestEntry,
-        local: bool = False,
-    ) -> Union[URIStr, FilePathStr]:
-        return self._handler.load_path(manifest_entry, local=local)
-
-    def store_path(
-        self,
-        artifact: "Artifact",
-        path: Union[URIStr, FilePathStr],
-        name: Optional[str] = None,
-        checksum: bool = True,
-        max_objects: Optional[int] = None,
-    ) -> Sequence[ArtifactManifestEntry]:
-        return self._handler.store_path(
-            artifact, path, name=name, checksum=checksum, max_objects=max_objects
-        )

wandb/sdk/interface/interface_grpc.py

@@ -1,460 +0,0 @@
-"""Backend Sender - Send to internal process.
-
-Manage backend sender.
-
-"""
-
-import logging
-from typing import TYPE_CHECKING, Any, Optional
-
-import grpc
-
-from wandb.proto import wandb_internal_pb2 as pb
-from wandb.proto import wandb_server_pb2_grpc as pbgrpc
-from wandb.proto import wandb_telemetry_pb2 as tpb
-
-from ..lib.mailbox import Mailbox, MailboxHandle
-from .interface import InterfaceBase
-from .message_future import MessageFuture
-from .message_future_poll import MessageFuturePoll
-
-if TYPE_CHECKING:
-    from ..wandb_run import Run
-
-
-logger = logging.getLogger("wandb")
-
-
-class InterfaceGrpc(InterfaceBase):
-    _stub: Optional[pbgrpc.InternalServiceStub]
-    _stream_id: Optional[str]
-    _mailbox: Mailbox
-
-    def __init__(self, mailbox: Mailbox) -> None:
-        super().__init__()
-        self._stub = None
-        self._process_check = None
-        self._stream_id = None
-        self._mailbox = mailbox
-
-    def _hack_set_run(self, run: "Run") -> None:
-        super()._hack_set_run(run)
-        assert run._run_id
-        self._stream_id = run._run_id
-
-    def _connect(self, stub: pbgrpc.InternalServiceStub) -> None:
-        self._stub = stub
-
-    def _assign(self, record: Any) -> None:
-        assert self._stream_id
-        record._info.stream_id = self._stream_id
-
-    def _communicate_check_version(
-        self, check_version: pb.CheckVersionRequest
-    ) -> Optional[pb.CheckVersionResponse]:
-        assert self._stub
-        self._assign(check_version)
-        run_result = self._stub.CheckVersion(check_version)
-        return run_result  # type: ignore
-
-    def _communicate_attach(
-        self, attach: pb.AttachRequest
-    ) -> Optional[pb.AttachResponse]:
-        assert self._stub
-        self._assign(attach)
-        resp = self._stub.Attach(attach)
-        return resp  # type: ignore
-
-    def _communicate_run(
-        self, run: pb.RunRecord, timeout: Optional[int] = None
-    ) -> Optional[pb.RunUpdateResult]:
-        assert self._stub
-        self._assign(run)
-        run_result = self._stub.RunUpdate(run)
-        return run_result  # type: ignore
-
-    def _publish_run(self, run: pb.RunRecord) -> None:
-        assert self._stub
-        self._assign(run)
-        _ = self._stub.RunUpdate(run)
-
-    def _publish_cancel(self, cancel: pb.CancelRequest) -> None:
-        assert self._stub
-        self._assign(cancel)
-        _ = self._stub.Cancel(cancel)
-
-    def _publish_config(self, cfg: pb.ConfigRecord) -> None:
-        assert self._stub
-        self._assign(cfg)
-        _ = self._stub.Config(cfg)
-
-    def _publish_metric(self, metric: pb.MetricRecord) -> None:
-        assert self._stub
-        self._assign(metric)
-        _ = self._stub.Metric(metric)
-
-    def _publish_summary(self, summary: pb.SummaryRecord) -> None:
-        assert self._stub
-        self._assign(summary)
-        _ = self._stub.Summary(summary)
-
-    def _communicate_get_summary(
-        self, get_summary: pb.GetSummaryRequest
-    ) -> Optional[pb.GetSummaryResponse]:
-        assert self._stub
-        self._assign(get_summary)
-        try:
-            resp = self._stub.GetSummary(get_summary)
-        except grpc.RpcError as e:
-            logger.info(f"GET SUMMARY TIMEOUT: {e}")
-            resp = pb.GetSummaryResponse()
-        return resp  # type: ignore
-
-    def _publish_telemetry(self, telem: tpb.TelemetryRecord) -> None:
-        assert self._stub
-        self._assign(telem)
-        _ = self._stub.Telemetry(telem)
-
-    def _publish_partial_history(
-        self, partial_history: pb.PartialHistoryRequest
-    ) -> None:
-        assert self._stub
-        self._assign(partial_history)
-        _ = self._stub.PartialLog(partial_history)
-
-    def _publish_history(self, history: pb.HistoryRecord) -> None:
-        assert self._stub
-        self._assign(history)
-        _ = self._stub.Log(history)
-
-    def _publish_preempting(self, preempt_rec: pb.RunPreemptingRecord) -> None:
-        assert self._stub
-        self._assign(preempt_rec)
-        _ = self._stub.RunPreempting(preempt_rec)
-
-    def _publish_output(self, outdata: pb.OutputRecord) -> None:
-        assert self._stub
-        self._assign(outdata)
-        _ = self._stub.Output(outdata)
-
-    def _publish_output_raw(self, outdata: pb.OutputRawRecord) -> None:
-        assert self._stub
-        self._assign(outdata)
-        _ = self._stub.OutputRaw(outdata)
-
-    def _communicate_shutdown(self) -> None:
-        assert self._stub
-        shutdown = pb.ShutdownRequest()
-        self._assign(shutdown)
-        _ = self._stub.Shutdown(shutdown)
-
-    def _communicate_run_start(
-        self, run_start: pb.RunStartRequest
-    ) -> Optional[pb.RunStartResponse]:
-        assert self._stub
-        self._assign(run_start)
-        try:
-            run_start_response = self._stub.RunStart(run_start)
-        except grpc.RpcError as e:
-            logger.info(f"RUNSTART TIMEOUT: {e}")
-            run_start_response = pb.RunStartResponse()
-        return run_start_response  # type: ignore
-
-    def _publish_files(self, files: pb.FilesRecord) -> None:
-        assert self._stub
-        self._assign(files)
-        _ = self._stub.Files(files)
-
-    def _publish_artifact(self, proto_artifact: pb.ArtifactRecord) -> None:
-        assert self._stub
-        self._assign(proto_artifact)
-        _ = self._stub.Artifact(proto_artifact)
-
-    def _publish_link_artifact(self, link_artifact: pb.LinkArtifactRecord) -> None:
-        assert self._stub
-        self._assign(link_artifact)
-        _ = self._stub.LinkArtifact(link_artifact)
-
-    def _publish_use_artifact(self, use_artifact: pb.UseArtifactRecord) -> None:
-        assert self._stub
-        self._assign(use_artifact)
-        _ = self._stub.UseArtifact(use_artifact)
-
-    def _communicate_artifact(
-        self, log_artifact: pb.LogArtifactRequest
-    ) -> MessageFuture:
-        art_send = pb.ArtifactSendRequest()
-        art_send.artifact.CopyFrom(log_artifact.artifact)
-
-        assert self._stub
-        self._assign(art_send)
-        art_send_resp = self._stub.ArtifactSend(art_send)
-
-        xid = art_send_resp.xid
-        future = MessageFuturePoll(self._future_poll_artifact, xid)
-        return future
-
-    def _future_poll_artifact(self, xid: str) -> Optional[pb.Result]:
-        art_poll = pb.ArtifactPollRequest(xid=xid)
-
-        assert self._stub
-        self._assign(art_poll)
-        art_poll_resp = self._stub.ArtifactPoll(art_poll)
-
-        if not art_poll_resp.ready:
-            return None
-
-        # emulate log_artifact response for old _communicate_artifact() protocol
-        result = pb.Result()
-        result.response.log_artifact_response.artifact_id = art_poll_resp.artifact_id
-        result.response.log_artifact_response.error_message = (
-            art_poll_resp.error_message
-        )
-        return result
-
-    def _communicate_artifact_send(
-        self, art_send: pb.ArtifactSendRequest
-    ) -> Optional[pb.ArtifactSendResponse]:
-        # NOTE: Not used in grpc interface now
-        raise NotImplementedError
-
-    def _communicate_artifact_poll(
-        self, art_send: pb.ArtifactPollRequest
-    ) -> Optional[pb.ArtifactPollResponse]:
-        # NOTE: Not used in grpc interface now
-        raise NotImplementedError
-
-    def _publish_artifact_done(self, artifact_done: pb.ArtifactDoneRequest) -> None:
-        # NOTE: Not used in grpc interface now
-        raise NotImplementedError
-
-    def _communicate_status(
-        self, status: pb.StatusRequest
-    ) -> Optional[pb.StatusResponse]:
-        assert self._stub
-        self._assign(status)
-        status_response = self._stub.Status(status)
-        return status_response  # type: ignore
-
-    def _communicate_network_status(
-        self, status: pb.NetworkStatusRequest
-    ) -> Optional[pb.NetworkStatusResponse]:
-        assert self._stub
-        self._assign(status)
-        # TODO: implement
-        return None
-
-    def _deliver_network_status(self, status: pb.NetworkStatusRequest) -> MailboxHandle:
-        assert self._stub
-        self._assign(status)
-        network_status_response = pb.NetworkStatusResponse()
-        response = pb.Response(network_status_response=network_status_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_stop_status(self, status: pb.StopStatusRequest) -> MailboxHandle:
-        assert self._stub
-        self._assign(status)
-        stop_status_response = pb.StopStatusResponse()
-        response = pb.Response(stop_status_response=stop_status_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _communicate_stop_status(
-        self, status: pb.StopStatusRequest
-    ) -> Optional[pb.StopStatusResponse]:
-        assert self._stub
-        self._assign(status)
-        # TODO: implement
-        return None
-
-    def _publish_alert(self, alert: pb.AlertRecord) -> None:
-        assert self._stub
-        self._assign(alert)
-        _ = self._stub.Alert(alert)
-
-    def _publish_tbdata(self, tbrecord: pb.TBRecord) -> None:
-        assert self._stub
-        self._assign(tbrecord)
-        _ = self._stub.TBSend(tbrecord)
-
-    def _publish_exit(self, exit_data: pb.RunExitRecord) -> None:
-        assert self._stub
-        self._assign(exit_data)
-        _ = self._stub.RunExit(exit_data)
-        return None
-
-    def _communicate_poll_exit(
-        self, poll_exit: pb.PollExitRequest
-    ) -> Optional[pb.PollExitResponse]:
-        assert self._stub
-        self._assign(poll_exit)
-        try:
-            ret = self._stub.PollExit(poll_exit)
-        except grpc.RpcError as e:
-            logger.info(f"POLL EXIT TIMEOUT: {e}")
-            ret = pb.PollExitResponse()
-        return ret  # type: ignore
-
-    def _publish_keepalive(self, keepalive: pb.KeepaliveRequest) -> None:
-        assert self._stub
-        self._assign(keepalive)
-        _ = self._stub.Keepalive(keepalive)
-        return None
-
-    def _communicate_server_info(
-        self, server_info: pb.ServerInfoRequest
-    ) -> Optional[pb.ServerInfoResponse]:
-        assert self._stub
-        self._assign(server_info)
-        ret = self._stub.ServerInfo(server_info)
-        return ret  # type: ignore
-
-    def _communicate_sampled_history(
-        self, sampled_history: pb.SampledHistoryRequest
-    ) -> Optional[pb.SampledHistoryResponse]:
-        assert self._stub
-        self._assign(sampled_history)
-        ret = self._stub.SampledHistory(sampled_history)
-        return ret  # type: ignore
-
-    def _publish_header(self, header: pb.HeaderRecord) -> None:
-        assert self._stub
-        # TODO: implement?
-
-    def _publish_pause(self, pause: pb.PauseRequest) -> None:
-        assert self._stub
-        self._assign(pause)
-        _ = self._stub.Pause(pause)
-
-    def _publish_resume(self, resume: pb.ResumeRequest) -> None:
-        assert self._stub
-        self._assign(resume)
-        _ = self._stub.Resume(resume)
-
-    def _deliver(self, result: pb.Result) -> MailboxHandle:
-        # TODO: rework grpc implementation to be asynchronous?
-        # for now we will just emulate the mailbox protocol
-        handle = self._mailbox.get_handle()
-        mailbox_slot = handle.address
-        result.control.mailbox_slot = mailbox_slot
-        self._mailbox.deliver(result)
-        return handle
-
-    def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle:
-        assert self._stub
-        self._assign(run)
-        run_result = self._stub.RunUpdate(run)
-        result = pb.Result(run_result=run_result)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_check_version(
-        self, check_version: pb.CheckVersionRequest
-    ) -> MailboxHandle:
-        assert self._stub
-        self._assign(check_version)
-        check_version_response = self._stub.CheckVersion(check_version)
-        response = pb.Response(check_version_response=check_version_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_attach(self, attach: pb.AttachRequest) -> MailboxHandle:
-        assert self._stub
-        self._assign(attach)
-        attach_response = self._stub.Attach(attach)
-        response = pb.Response(attach_response=attach_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_get_summary(self, get_summary: pb.GetSummaryRequest) -> MailboxHandle:
-        assert self._stub
-        self._assign(get_summary)
-        try:
-            get_summary_response = self._stub.GetSummary(get_summary)
-        except grpc.RpcError as e:
-            logger.info(f"GET SUMMARY TIMEOUT: {e}")
-            get_summary_response = pb.GetSummaryResponse()
-        response = pb.Response(get_summary_response=get_summary_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_exit(self, run_exit: pb.RunExitRecord) -> MailboxHandle:
-        assert self._stub
-        self._assign(run_exit)
-        exit_result = self._stub.RunExit(run_exit)
-        result = pb.Result(exit_result=exit_result)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_poll_exit(self, poll_exit: pb.PollExitRequest) -> MailboxHandle:
-        assert self._stub
-        self._assign(poll_exit)
-        poll_exit_response = self._stub.PollExit(poll_exit)
-        response = pb.Response(poll_exit_response=poll_exit_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_request_server_info(
-        self, server_info: pb.ServerInfoRequest
-    ) -> MailboxHandle:
-        assert self._stub
-        self._assign(server_info)
-        server_info_response = self._stub.ServerInfo(server_info)
-        response = pb.Response(server_info_response=server_info_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_request_sampled_history(
-        self, sampled_history: pb.SampledHistoryRequest
-    ) -> MailboxHandle:
-        assert self._stub
-        self._assign(sampled_history)
-        sampled_history_response = self._stub.SampledHistory(sampled_history)
-        response = pb.Response(sampled_history_response=sampled_history_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_request_run_status(
-        self, run_status: pb.RunStatusRequest
-    ) -> MailboxHandle:
-        assert self._stub
-        self._assign(run_status)
-        run_status_response = self._stub.RunStatus(run_status)
-        response = pb.Response(run_status_response=run_status_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_run_start(self, run_start: pb.RunStartRequest) -> MailboxHandle:
-        assert self._stub
-        self._assign(run_start)
-        try:
-            run_start_response = self._stub.RunStart(run_start)
-        except grpc.RpcError as e:
-            logger.info(f"RUNSTART TIMEOUT: {e}")
-            run_start_response = pb.RunStartResponse()
-        response = pb.Response(run_start_response=run_start_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def _deliver_request_job_info(self, job_info: pb.JobInfoRequest) -> MailboxHandle:
-        assert self._stub
-        self._assign(job_info)
-        job_info_response = self._stub.JobInfo(job_info)
-        response = pb.Response(job_info_response=job_info_response)
-        result = pb.Result(response=response)
-        handle = self._deliver(result)
-        return handle
-
-    def join(self) -> None:
-        super().join()