wandb 0.17.0rc2__py3-none-any.whl → 0.17.1__py3-none-any.whl

wandb/sdk/internal/file_stream.py

@@ -1,4 +1,3 @@
-import base64
 import functools
 import itertools
 import json
@@ -53,7 +52,7 @@ logger = logging.getLogger(__name__)
 
 class Chunk(NamedTuple):
     filename: str
-    data: Any
+    data: str
 
 
 class DefaultFilePolicy:
@@ -227,7 +226,7 @@ class CRDedupeFilePolicy(DefaultFilePolicy):
         prefix += token + " "
         return prefix, rest
 
-    def process_chunks(self, chunks: List) -> List["ProcessedChunk"]:
+    def process_chunks(self, chunks: List[Chunk]) -> List["ProcessedChunk"]:
         r"""Process chunks.
 
         Args:
@@ -300,18 +299,6 @@ class CRDedupeFilePolicy(DefaultFilePolicy):
         return ret
 
 
-class BinaryFilePolicy(DefaultFilePolicy):
-    def __init__(self) -> None:
-        super().__init__()
-        self._offset: int = 0
-
-    def process_chunks(self, chunks: List[Chunk]) -> "ProcessedBinaryChunk":
-        data = b"".join([c.data for c in chunks])
-        enc = base64.b64encode(data).decode("ascii")
-        self._offset += len(data)
-        return {"offset": self._offset, "content": enc, "encoding": "base64"}
-
-
 class FileStreamApi:
     """Pushes chunks of files to our streaming endpoint.
 
@@ -520,7 +507,7 @@ class FileStreamApi:
             wandb.termerror(
                 "Dropped streaming file chunk (see wandb/debug-internal.log)"
             )
-            logger.exception("dropped chunk %s" % response)
+            logger.exception("dropped chunk {}".format(response))
             self._dropped_chunks += 1
         else:
             parsed: Optional[dict] = None
@@ -585,12 +572,12 @@ class FileStreamApi:
     def enqueue_preempting(self) -> None:
         self._queue.put(self.Preempting())
 
-    def push(self, filename: str, data: Any) -> None:
+    def push(self, filename: str, data: str) -> None:
         """Push a chunk of a file to the streaming endpoint.
 
         Arguments:
-            filename: Name of file that this is a chunk of.
-            data: File data.
+            filename: Name of file to append to.
+            data: Text to append to the file.
         """
         self._queue.put(Chunk(filename, data))
 

wandb/sdk/internal/internal_api.py

@@ -1,5 +1,4 @@
 import ast
-import asyncio
 import base64
 import datetime
 import functools
@@ -49,7 +48,7 @@ from ..lib import retry
 from ..lib.filenames import DIFF_FNAME, METADATA_FNAME
 from ..lib.gitlib import GitRepo
 from . import context
-from .progress import AsyncProgress, Progress
+from .progress import Progress
 
 logger = logging.getLogger(__name__)
 
@@ -121,13 +120,6 @@ if TYPE_CHECKING:
     SweepState = Literal["RUNNING", "PAUSED", "CANCELED", "FINISHED"]
     Number = Union[int, float]
 
-# This funny if/else construction is the simplest thing I've found that
-# works at runtime, satisfies Mypy, and gives autocomplete in VSCode:
-if TYPE_CHECKING:
-    import httpx
-else:
-    httpx = util.get_module("httpx")
-
 # class _MappingSupportsCopy(Protocol):
 #     def copy(self) -> "_MappingSupportsCopy": ...
 #     def keys(self) -> Iterable: ...
@@ -161,23 +153,6 @@ def check_httpclient_logger_handler() -> None:
         httpclient_logger.addHandler(root_logger.handlers[0])
 
 
-def check_httpx_exc_retriable(exc: Exception) -> bool:
-    retriable_codes = (308, 408, 409, 429, 500, 502, 503, 504)
-    return (
-        isinstance(exc, (httpx.TimeoutException, httpx.NetworkError))
-        or (
-            isinstance(exc, httpx.HTTPStatusError)
-            and exc.response.status_code in retriable_codes
-        )
-        or (
-            isinstance(exc, httpx.HTTPStatusError)
-            and exc.response.status_code == 400
-            and "x-amz-meta-md5" in exc.request.headers
-            and "RequestTimeout" in str(exc.response.content)
-        )
-    )
-
-
 class _ThreadLocalData(threading.local):
     context: Optional[context.Context]
 
@@ -286,10 +261,6 @@ class Api:
             )
         )
 
-        # httpx is an optional dependency, so we lazily instantiate the client
-        # only when we need it
-        self._async_httpx_client: Optional[httpx.AsyncClient] = None
-
         self.retry_callback = retry_callback
         self._retry_gql = retry.Retry(
             self.execute,
@@ -361,7 +332,7 @@ class Api:
 
     def relocate(self) -> None:
         """Ensure the current api points to the right server."""
-        self.client.transport.url = "%s/graphql" % self.settings("base_url")
+        self.client.transport.url = "{}/graphql".format(self.settings("base_url"))
 
     def execute(self, *args: Any, **kwargs: Any) -> "_Response":
         """Wrapper around execute that logs in cases of failure."""
@@ -2245,6 +2216,113 @@ class Api:
             server_messages,
         )
 
+    @normalize_exceptions
+    def rewind_run(
+        self,
+        run_name: str,
+        metric_name: str,
+        metric_value: float,
+        program_path: Optional[str] = None,
+        entity: Optional[str] = None,
+        project: Optional[str] = None,
+        num_retries: Optional[int] = None,
+    ) -> dict:
+        """Rewinds a run to a previous state.
+
+        Arguments:
+            run_name (str): The name of the run to rewind
+            metric_name (str): The name of the metric to rewind to
+            metric_value (float): The value of the metric to rewind to
+            program_path (str, optional): Path to the program
+            entity (str, optional): The entity to scope this project to
+            project (str, optional): The name of the project
+            num_retries (int, optional): Number of retries
+
+        Returns:
+            A dict with the rewound run
+
+                {
+                    "id": "run_id",
+                    "name": "run_name",
+                    "displayName": "run_display_name",
+                    "description": "run_description",
+                    "config": "stringified_run_config_json",
+                    "sweepName": "run_sweep_name",
+                    "project": {
+                        "id": "project_id",
+                        "name": "project_name",
+                        "entity": {
+                            "id": "entity_id",
+                            "name": "entity_name"
+                        }
+                    },
+                    "historyLineCount": 100,
+                }
+        """
+        query_string = """
+        mutation RewindRun($runName: String!, $entity: String, $project: String, $metricName: String!, $metricValue: Float!) {
+            rewindRun(input: {runName: $runName, entityName: $entity, projectName: $project, metricName: $metricName, metricValue: $metricValue}) {
+                rewoundRun {
+                    id
+                    name
+                    displayName
+                    description
+                    config
+                    sweepName
+                    project {
+                        id
+                        name
+                        entity {
+                            id
+                            name
+                        }
+                    }
+                    historyLineCount
+                }
+            }
+        }
+        """
+
+        mutation = gql(query_string)
+
+        kwargs = {}
+        if num_retries is not None:
+            kwargs["num_retries"] = num_retries
+
+        variable_values = {
+            "runName": run_name,
+            "entity": entity or self.settings("entity"),
+            "project": project or util.auto_project_name(program_path),
+            "metricName": metric_name,
+            "metricValue": metric_value,
+        }
+
+        # retry conflict errors for 2 minutes, default to no_auth_retry
+        check_retry_fn = util.make_check_retry_fn(
+            check_fn=util.check_retry_conflict_or_gone,
+            check_timedelta=datetime.timedelta(minutes=2),
+            fallback_retry_fn=util.no_retry_auth,
+        )
+
+        response = self.gql(
+            mutation,
+            variable_values=variable_values,
+            check_retry_fn=check_retry_fn,
+            **kwargs,
+        )
+
+        run_obj: Dict[str, Dict[str, Dict[str, str]]] = response.get(
+            "rewindRun", {}
+        ).get("rewoundRun", {})
+        project_obj: Dict[str, Dict[str, str]] = run_obj.get("project", {})
+        if project_obj:
+            self.set_setting("project", project_obj["name"])
+            entity_obj = project_obj.get("entity", {})
+            if entity_obj:
+                self.set_setting("entity", entity_obj["name"])
+
+        return run_obj
+
     @normalize_exceptions
     def get_run_info(
         self,
@@ -2794,105 +2872,6 @@ class Api:
 
         return response
 
-    async def upload_file_async(
-        self,
-        url: str,
-        file: IO[bytes],
-        callback: Optional["ProgressFn"] = None,
-        extra_headers: Optional[Dict[str, str]] = None,
-    ) -> None:
-        """An async not-quite-equivalent version of `upload_file`.
-
-        Differences from `upload_file`:
-            - This method doesn't implement Azure uploads. (The Azure SDK supports
-              async, but it's nontrivial to use it here.) If the upload looks like
-              it's destined for Azure, this method will delegate to the sync impl.
-            - Consequently, this method doesn't return the response object.
-              (Because it might fall back to the sync impl, it would sometimes
-               return a `requests.Response` and sometimes an `httpx.Response`.)
-            - This method doesn't wrap retryable errors in `TransientError`.
-              It leaves that determination to the caller.
-        """
-        check_httpclient_logger_handler()
-        must_delegate = False
-
-        if httpx is None:
-            wandb.termwarn(  # type: ignore[unreachable]
-                "async file-uploads require `pip install wandb[async]`; falling back to sync implementation",
-                repeat=False,
-            )
-            must_delegate = True
-
-        if extra_headers is not None and "x-ms-blob-type" in extra_headers:
-            wandb.termwarn(
-                "async file-uploads don't support Azure; falling back to sync implementation",
-                repeat=False,
-            )
-            must_delegate = True
-
-        if must_delegate:
-            await asyncio.get_event_loop().run_in_executor(
-                None,
-                lambda: self.upload_file_retry(
-                    url=url,
-                    file=file,
-                    callback=callback,
-                    extra_headers=extra_headers,
-                ),
-            )
-            return
-
-        if self._async_httpx_client is None:
-            self._async_httpx_client = httpx.AsyncClient()
-
-        progress = AsyncProgress(Progress(file, callback=callback))
-
-        try:
-            response = await self._async_httpx_client.put(
-                url=url,
-                content=progress,
-                headers={
-                    "Content-Length": str(len(progress)),
-                    **(extra_headers if extra_headers is not None else {}),
-                },
-            )
-            response.raise_for_status()
-        except Exception as e:
-            progress.rewind()
-            logger.error(f"upload_file_async exception {url}: {e}")
-            if isinstance(e, httpx.RequestError):
-                logger.error(f"upload_file_async request headers: {e.request.headers}")
-            if isinstance(e, httpx.HTTPStatusError):
-                logger.error(f"upload_file_async response body: {e.response.content!r}")
-            raise
-
-    async def upload_file_retry_async(
-        self,
-        url: str,
-        file: IO[bytes],
-        callback: Optional["ProgressFn"] = None,
-        extra_headers: Optional[Dict[str, str]] = None,
-        num_retries: int = 100,
-    ) -> None:
-        backoff = retry.FilteredBackoff(
-            filter=check_httpx_exc_retriable,
-            wrapped=retry.ExponentialBackoff(
-                initial_sleep=datetime.timedelta(seconds=1),
-                max_sleep=datetime.timedelta(seconds=60),
-                max_retries=num_retries,
-                timeout_at=datetime.datetime.now() + datetime.timedelta(days=7),
-            ),
-        )
-
-        await retry.retry_async(
-            backoff=backoff,
-            fn=self.upload_file_async,
-            url=url,
-            file=file,
-            callback=callback,
-            extra_headers=extra_headers,
-        )
-
     @normalize_exceptions
     def register_agent(
         self,
@@ -3039,9 +3018,10 @@ class Api:
                         parameter["distribution"] = "uniform"
                     else:
                         raise ValueError(
-                            "Parameter %s is ambiguous, please specify bounds as both floats (for a float_"
-                            "uniform distribution) or ints (for an int_uniform distribution)."
-                            % parameter_name
+                            "Parameter {} is ambiguous, please specify bounds as both floats (for a float_"
+                            "uniform distribution) or ints (for an int_uniform distribution).".format(
+                                parameter_name
+                            )
                         )
         return config
 
@@ -3144,7 +3124,9 @@ class Api:
 
         # Silly, but attr-dicts like EasyDicts don't serialize correctly to yaml.
         # This sanitizes them with a round trip pass through json to get a regular dict.
-        config_str = yaml.dump(json.loads(json.dumps(config)))
+        config_str = yaml.dump(
+            json.loads(json.dumps(config)), Dumper=util.NonOctalStringDumper
+        )
 
         err: Optional[Exception] = None
         for mutation in mutations:
@@ -3887,6 +3869,36 @@ class Api:
             response["updateArtifactManifest"]["artifactManifest"]["file"],
         )
 
+    def update_artifact_metadata(
+        self, artifact_id: str, metadata: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Set the metadata of the given artifact version."""
+        mutation = gql(
+            """
+        mutation UpdateArtifact(
+            $artifactID: ID!,
+            $metadata: JSONString,
+        ) {
+            updateArtifact(input: {
+                artifactID: $artifactID,
+                metadata: $metadata,
+            }) {
+                artifact {
+                    id
+                }
+            }
+        }
+        """
+        )
+        response = self.gql(
+            mutation,
+            variable_values={
+                "artifactID": artifact_id,
+                "metadata": json.dumps(metadata),
+            },
+        )
+        return response["updateArtifact"]["artifact"]
+
     def _resolve_client_id(
         self,
         client_id: str,
@@ -4081,9 +4093,9 @@ class Api:
         s = self.sweep(sweep=sweep, entity=entity, project=project, specs="{}")
         curr_state = s["state"].upper()
         if state == "PAUSED" and curr_state not in ("PAUSED", "RUNNING"):
-            raise Exception("Cannot pause %s sweep." % curr_state.lower())
+            raise Exception("Cannot pause {} sweep.".format(curr_state.lower()))
         elif state != "RUNNING" and curr_state not in ("RUNNING", "PAUSED", "PENDING"):
-            raise Exception("Sweep already %s." % curr_state.lower())
+            raise Exception("Sweep already {}.".format(curr_state.lower()))
         sweep_id = s["id"]
         mutation = gql(
             """