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

wandb/sdk/interface/interface_shared.py

@@ -100,6 +100,10 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(telemetry=telem)
         self._publish(rec)
 
+    def _publish_job_input(self, job_input: pb.JobInputRequest) -> MailboxHandle:
+        record = self._make_request(job_input=job_input)
+        return self._deliver_record(record)
+
     def _make_stats(self, stats_dict: dict) -> pb.StatsRecord:
         stats = pb.StatsRecord()
         stats.stats_type = pb.StatsRecord.StatsType.SYSTEM
@@ -147,6 +151,7 @@ class InterfaceShared(InterfaceBase):
         telemetry_record: Optional[pb.TelemetryRecordRequest] = None,
         get_system_metrics: Optional[pb.GetSystemMetricsRequest] = None,
         python_packages: Optional[pb.PythonPackagesRequest] = None,
+        job_input: Optional[pb.JobInputRequest] = None,
     ) -> pb.Record:
         request = pb.Request()
         if login:
@@ -207,6 +212,8 @@ class InterfaceShared(InterfaceBase):
             request.sync.CopyFrom(sync)
         elif python_packages:
             request.python_packages.CopyFrom(python_packages)
+        elif job_input:
+            request.job_input.CopyFrom(job_input)
         else:
             raise Exception("Invalid request")
         record = self._make_record(request=request)
@@ -239,9 +246,6 @@ class InterfaceShared(InterfaceBase):
         use_artifact: Optional[pb.UseArtifactRecord] = None,
         output: Optional[pb.OutputRecord] = None,
         output_raw: Optional[pb.OutputRawRecord] = None,
-        launch_wandb_config_parameters: Optional[
-            pb.LaunchWandbConfigParametersRecord
-        ] = None,
     ) -> pb.Record:
         record = pb.Record()
         if run:
@@ -286,8 +290,6 @@ class InterfaceShared(InterfaceBase):
             record.output.CopyFrom(output)
         elif output_raw:
             record.output_raw.CopyFrom(output_raw)
-        elif launch_wandb_config_parameters:
-            record.wandb_config_parameters.CopyFrom(launch_wandb_config_parameters)
         else:
             raise Exception("Invalid record")
         return record
@@ -417,14 +419,6 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(alert=proto_alert)
         self._publish(rec)
 
-    def _publish_launch_wandb_config_parameters(
-        self, launch_wandb_config_parameters: pb.LaunchWandbConfigParametersRecord
-    ) -> None:
-        rec = self._make_record(
-            launch_wandb_config_parameters=launch_wandb_config_parameters
-        )
-        self._publish(rec)
-
     def _communicate_status(
         self, status: pb.StatusRequest
     ) -> Optional[pb.StatusResponse]:

wandb/sdk/internal/datastore.py

@@ -52,7 +52,7 @@ try:
     bytes("", "ascii")
 
     def strtobytes(x):
-        """strtobytes."""
+        """Strtobytes."""
         return bytes(x, "iso8859-1")
 
     # def bytestostr(x):

wandb/sdk/internal/file_pusher.py

@@ -14,7 +14,7 @@ from wandb.sdk.lib.paths import LogicalPath
 
 if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
-    from wandb.sdk.artifacts.artifact_saver import SaveFn, SaveFnAsync
+    from wandb.sdk.artifacts.artifact_saver import SaveFn
     from wandb.sdk.internal import file_stream, internal_api
     from wandb.sdk.internal.settings_static import SettingsStatic
 
@@ -148,11 +148,8 @@ class FilePusher:
         manifest: "ArtifactManifest",
         artifact_id: str,
         save_fn: "SaveFn",
-        save_fn_async: "SaveFnAsync",
     ) -> None:
-        event = step_checksum.RequestStoreManifestFiles(
-            manifest, artifact_id, save_fn, save_fn_async
-        )
+        event = step_checksum.RequestStoreManifestFiles(manifest, artifact_id, save_fn)
         self._incoming_queue.put(event)
 
     def commit_artifact(

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.
 
@@ -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/handler.py

@@ -50,6 +50,18 @@ SummaryDict = Dict[str, Any]
 
 logger = logging.getLogger(__name__)
 
+# Update (March 5, 2024): Since ~2020/2021, when constructing the summary
+# object, we had replaced the artifact path for media types with the latest
+# artifact path. The primary purpose of this was to support live updating of
+# media objects in the UI (since the default artifact path was fully qualified
+# and would not update). However, in March of 2024, a bug was discovered with
+# this approach which causes this path to be incorrect in cases where the media
+# object is logged to another artifact before being logged to the run. Setting
+# this to `False` disables this copy behavior. The impact is that users will
+# need to refresh to see updates. Ironically, this updating behavior is not
+# currently supported in the UI, so the impact of this change is minimal.
+REPLACE_SUMMARY_ART_PATH_WITH_LATEST = False
+
 
 def _dict_nested_set(target: Dict[str, Any], key_list: Sequence[str], v: Any) -> None:
     # recurse down the dictionary structure:
@@ -371,7 +383,11 @@ class HandleManager:
                     updated = True
             return updated
         # If the dict is a media object, update the pointer to the latest alias
-        elif isinstance(v, dict) and handler_util.metric_is_wandb_dict(v):
+        elif (
+            REPLACE_SUMMARY_ART_PATH_WITH_LATEST
+            and isinstance(v, dict)
+            and handler_util.metric_is_wandb_dict(v)
+        ):
             if "_latest_artifact_path" in v and "artifact_path" in v:
                 # TODO: Make non-destructive?
                 v["artifact_path"] = v["_latest_artifact_path"]
@@ -381,7 +397,7 @@ class HandleManager:
     def _update_summary_media_objects(self, v: Dict[str, Any]) -> Dict[str, Any]:
         # For now, non-recursive - just top level
         for nk, nv in v.items():
-            if (
+            if REPLACE_SUMMARY_ART_PATH_WITH_LATEST and (
                 isinstance(nv, dict)
                 and handler_util.metric_is_wandb_dict(nv)
                 and "_latest_artifact_path" in nv

wandb/sdk/internal/internal.py

@@ -12,7 +12,6 @@ Threads:
 
 """
 
-
 import atexit
 import logging
 import os

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,
@@ -2794,105 +2765,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,

wandb/sdk/internal/internal_util.py

@@ -4,7 +4,6 @@ Collection of classes to support the internal process.
 
 """
 
-
 import logging
 import queue
 import sys

wandb/sdk/internal/job_builder.py

@@ -1,4 +1,5 @@
 """job builder."""
+
 import json
 import logging
 import os
@@ -31,6 +32,43 @@ JOB_ARTIFACT_TYPE = "job"
 LOG_LEVEL = Literal["log", "warn", "error"]
 
 
+class Version:
+    def __init__(self, major: int, minor: int, patch: int):
+        self._major = major
+        self._minor = minor
+        self._patch = patch
+
+    def __repr__(self) -> str:
+        return f"{self._major}.{self._minor}.{self._patch}"
+
+    def __lt__(self, other: "Version") -> bool:
+        if self._major < other._major:
+            return True
+        elif self._major == other._major:
+            if self._minor < other._minor:
+                return True
+            elif self._minor == other._minor:
+                if self._patch < other._patch:
+                    return True
+        return False
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Version):
+            return NotImplemented
+        return (
+            self._major == other._major
+            and self._minor == other._minor
+            and self._patch == other._patch
+        )
+
+
+# Minimum supported wandb version for keys in the source dict of wandb-job.json
+SOURCE_KEYS_MIN_SUPPORTED_VERSION = {
+    "dockerfile": Version(0, 17, 0),
+    "build_context": Version(0, 17, 0),
+}
+
+
 class GitInfo(TypedDict):
     remote: str
     commit: str
@@ -40,12 +78,16 @@ class GitSourceDict(TypedDict):
     git: GitInfo
     entrypoint: List[str]
     notebook: bool
+    build_context: Optional[str]
+    dockerfile: Optional[str]
 
 
 class ArtifactSourceDict(TypedDict):
     artifact: str
     entrypoint: List[str]
     notebook: bool
+    build_context: Optional[str]
+    dockerfile: Optional[str]
 
 
 class ImageSourceDict(TypedDict):
@@ -72,6 +114,19 @@ class ArtifactInfoForJob(TypedDict):
     name: str
 
 
+def get_min_supported_for_source_dict(
+    source: Union[GitSourceDict, ArtifactSourceDict, ImageSourceDict],
+) -> Optional[Version]:
+    """Get the minimum supported wandb version the source dict of wandb-job.json."""
+    min_seen = None
+    for key in source:
+        new_ver = SOURCE_KEYS_MIN_SUPPORTED_VERSION.get(key)
+        if new_ver:
+            if min_seen is None or new_ver < min_seen:
+                min_seen = new_ver
+    return min_seen
+
+
 class JobArtifact(Artifact):
     def __init__(self, name: str, *args: Any, **kwargs: Any):
         super().__init__(name, "placeholder", *args, **kwargs)
@@ -105,9 +160,9 @@ class JobBuilder:
         self._disable = settings.disable_job_creation
         self._partial_source = None
         self._aliases = []
-        self._source_type: Optional[
-            Literal["repo", "artifact", "image"]
-        ] = settings.job_source  # type: ignore[assignment]
+        self._source_type: Optional[Literal["repo", "artifact", "image"]] = (
+            settings.job_source  # type: ignore[assignment]
+        )
         self._is_notebook_run = self._get_is_notebook_run()
         self._verbose = verbose
 
@@ -196,6 +251,8 @@ class JobBuilder:
             "git": {"remote": remote, "commit": commit},
             "entrypoint": entrypoint,
             "notebook": self._is_notebook_run,
+            "build_context": metadata.get("build_context"),
+            "dockerfile": metadata.get("dockerfile"),
         }
         name = self._make_job_name(f"{remote}_{program_relpath}")
 
@@ -246,6 +303,8 @@ class JobBuilder:
             "entrypoint": entrypoint,
             "notebook": self._is_notebook_run,
             "artifact": f"wandb-artifact://_id/{self._logged_code_artifact['id']}",
+            "build_context": metadata.get("build_context"),
+            "dockerfile": metadata.get("dockerfile"),
         }
         name = self._make_job_name(self._logged_code_artifact["name"])
 
@@ -275,7 +334,7 @@ class JobBuilder:
         return source, name
 
     def _make_job_name(self, input_str: str) -> str:
-        """Use job name from settings if provided, else use programatic name."""
+        """Use job name from settings if provided, else use programmatic name."""
         if self._settings.job_name:
             return self._settings.job_name
 
@@ -292,17 +351,6 @@ class JobBuilder:
             if metadata.get("entrypoint"):
                 entrypoint: List[str] = metadata["entrypoint"]
                 return entrypoint
-
-            # if entrypoint is not in metadata, then construct from python
-            assert metadata.get("python")
-
-            python = metadata["python"]
-            if python.count(".") > 1:
-                python = ".".join(python.split(".")[:2])
-
-            entrypoint = [f"python{python}", program_relpath]
-            return entrypoint
-
         # job is being built from a run
         entrypoint = [os.path.basename(sys.executable), program_relpath]
 
@@ -314,7 +362,73 @@ class JobBuilder:
     def _is_colab_run(self) -> bool:
         return hasattr(self._settings, "_colab") and bool(self._settings._colab)
 
-    def build(self) -> Optional[Artifact]:
+    def _build_job_source(
+        self,
+        source_type: str,
+        program_relpath: Optional[str],
+        metadata: Dict[str, Any],
+    ) -> Tuple[
+        Union[GitSourceDict, ArtifactSourceDict, ImageSourceDict, None],
+        Optional[str],
+    ]:
+        """Construct a job source dict and name from the current run.
+
+        Arguments:
+            source_type (str): The type of source to build the job from. One of
+                "repo", "artifact", or "image".
+        """
+        source: Union[
+            GitSourceDict,
+            ArtifactSourceDict,
+            ImageSourceDict,
+            None,
+        ] = None
+
+        if source_type == "repo":
+            source, name = self._build_repo_job_source(
+                program_relpath or "",
+                metadata,
+            )
+        elif source_type == "artifact":
+            source, name = self._build_artifact_job_source(
+                program_relpath or "",
+                metadata,
+            )
+        elif source_type == "image" and self._has_image_job_ingredients(metadata):
+            source, name = self._build_image_job_source(metadata)
+        else:
+            source = None
+
+        if source is None:
+            if source_type:
+                self._log_if_verbose(
+                    f"Source type is set to '{source_type}' but some required information is missing "
+                    "from the environment. A job will not be created from this run. See "
+                    "https://docs.wandb.ai/guides/launch/create-job",
+                    "warn",
+                )
+            return None, None
+
+        return source, name
+
+    def build(
+        self,
+        build_context: Optional[str] = None,
+        dockerfile: Optional[str] = None,
+    ) -> Optional[Artifact]:
+        """Build a job artifact from the current run.
+
+        Arguments:
+            build_context (Optional[str]): Path within the job source code to
+                the image build context. Saved as part of the job for future
+                builds.
+            dockerfile (Optional[str]): Path within the build context the
+                Dockerfile. Saved as part of the job for future builds.
+
+        Returns:
+            Optional[Artifact]: The job artifact if it was successfully built,
+            otherwise None.
+        """
         _logger.info("Attempting to build job artifact")
         if not os.path.exists(
             os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME)
@@ -331,6 +445,8 @@ class JobBuilder:
                 "warn",
             )
             return None
+        metadata["dockerfile"] = dockerfile
+        metadata["build_context"] = build_context
 
         runtime: Optional[str] = metadata.get("python")
         # can't build a job without a python version
@@ -374,45 +490,39 @@ class JobBuilder:
                 return None
 
             program_relpath = self._get_program_relpath(source_type, metadata)
-            if source_type != "image" and not program_relpath:
+            if (
+                not metadata.get("_partial")
+                and source_type != "image"
+                and not program_relpath
+            ):
                 self._log_if_verbose(
                     "No program path found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
                     "warn",
                 )
                 return None
 
-            source: Union[
-                Optional[GitSourceDict],
-                Optional[ArtifactSourceDict],
-                Optional[ImageSourceDict],
-            ] = None
-
-            # make source dict
-            if source_type == "repo":
-                assert program_relpath
-                source, name = self._build_repo_job_source(program_relpath, metadata)
-            elif source_type == "artifact":
-                assert program_relpath
-                source, name = self._build_artifact_job_source(
-                    program_relpath, metadata
-                )
-            elif source_type == "image" and self._has_image_job_ingredients(metadata):
-                source, name = self._build_image_job_source(metadata)
-            else:
-                source = None
-
+            source, name = self._build_job_source(
+                source_type,
+                program_relpath,
+                metadata,
+            )
             if source is None:
-                if source_type:
-                    self._log_if_verbose(
-                        f"Source type is set to '{source_type}' but some required information is missing "
-                        "from the environment. A job will not be created from this run. See "
-                        "https://docs.wandb.ai/guides/launch/create-job",
-                        "warn",
-                    )
                 return None
 
+            if build_context:
+                source["build_context"] = build_context  # type: ignore[typeddict-item]
+            if dockerfile:
+                source["dockerfile"] = dockerfile  # type: ignore[typeddict-item]
+
+            # Pop any keys that are initialized to None. The current TypedDict
+            # system for source dicts requires all keys to be present, but we
+            # don't want to include keys that are None in the final dict.
+            for key in list(source.keys()):
+                if source[key] is None:  # type: ignore[literal-required]
+                    source.pop(key)  # type: ignore[literal-require,misc]
+
             source_info = {
-                "_version": "v0",
+                "_version": str(get_min_supported_for_source_dict(source) or "v0"),
                 "source_type": source_type,
                 "source": source,
                 "input_types": input_types,
@@ -530,6 +640,8 @@ def convert_use_artifact_to_job_source(
             },
             "entrypoint": entrypoint,
             "notebook": source_info.source.git.notebook,
+            "build_context": None,
+            "dockerfile": None,
         }
         source_info_dict.update({"source": git_source})
     elif source_info.source_type == "artifact":
@@ -538,6 +650,8 @@ def convert_use_artifact_to_job_source(
             "artifact": source_info.source.artifact.artifact,
             "entrypoint": entrypoint,
             "notebook": source_info.source.artifact.notebook,
+            "build_context": None,
+            "dockerfile": None,
         }
         source_info_dict.update({"source": artifact_source})
     elif source_info.source_type == "image":