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

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":

wandb/sdk/internal/profiler.py

@@ -1,4 +1,5 @@
 """Integration with pytorch profiler."""
+
 import os
 
 import wandb

wandb/sdk/internal/progress.py

@@ -81,31 +81,3 @@ class Progress:
         return self.len
 
     next = __next__
-
-
-class AsyncProgress:
-    """Wrapper around Progress, to make it async iterable.
-
-    httpx, for streaming uploads, requires the data source to be an async iterable.
-    If we pass in a sync iterable (like a bare `Progress` instance), httpx will
-    get confused, think we're trying to make a synchronous request, and raise.
-    So we need this wrapper class to be an async iterable but *not* a sync iterable.
-    """
-
-    def __init__(self, progress: Progress) -> None:
-        self._progress = progress
-
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self):
-        try:
-            return next(self._progress)
-        except StopIteration:
-            raise StopAsyncIteration
-
-    def __len__(self):
-        return len(self._progress)
-
-    def rewind(self) -> None:
-        self._progress.rewind()

wandb/sdk/internal/run.py

@@ -4,6 +4,7 @@
 Semi-stubbed run for internal process use.
 
 """
+
 from wandb._globals import _datatypes_set_callback
 
 from .. import wandb_run

wandb/sdk/internal/sender.py

@@ -327,7 +327,6 @@ class SendManager:
             # ignore_globs=(),
             _sync=True,
             disable_job_creation=False,
-            _async_upload_concurrency_limit=None,
             _file_stream_timeout_seconds=0,
         )
         record_q: Queue[Record] = queue.Queue()
@@ -910,7 +909,7 @@ class SendManager:
         is_wandb_init = self._run is None
 
         # save start time of a run
-        self._start_time = run.start_time.ToMicroseconds() // 1e6
+        self._start_time = int(run.start_time.ToMicroseconds() // 1e6)
 
         # update telemetry
         if run.telemetry:

wandb/sdk/internal/system/assets/gpu_amd.py

@@ -28,14 +28,6 @@ logger = logging.getLogger(__name__)
 ROCM_SMI_CMD: Final[str] = shutil.which("rocm-smi") or "/usr/bin/rocm-smi"
 
 
-def get_rocm_smi_stats() -> Dict[str, Any]:
-    command = [str(ROCM_SMI_CMD), "-a", "--json"]
-    output = subprocess.check_output(command, universal_newlines=True).strip()
-    if "No AMD GPUs specified" in output:
-        return {}
-    return json.loads(output.split("\n")[0])  # type: ignore
-
-
 _StatsKeys = Literal[
     "gpu",
     "memoryAllocated",
@@ -49,6 +41,48 @@ _Stats = Dict[_StatsKeys, float]
 _InfoDict = Dict[str, Union[int, List[Dict[str, Any]]]]
 
 
+def get_rocm_smi_stats() -> Dict[str, Any]:
+    command = [str(ROCM_SMI_CMD), "-a", "--json"]
+    output = subprocess.check_output(command, universal_newlines=True).strip()
+    if "No AMD GPUs specified" in output:
+        return {}
+    return json.loads(output.split("\n")[0])  # type: ignore
+
+
+def parse_stats(stats: Dict[str, str]) -> _Stats:
+    """Parse stats from rocm-smi output."""
+    parsed_stats: _Stats = {}
+
+    try:
+        parsed_stats["gpu"] = float(stats.get("GPU use (%)"))  # type: ignore
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU usage as float")
+    try:
+        parsed_stats["memoryAllocated"] = float(stats.get("GPU memory use (%)"))  # type: ignore
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU memory allocation as float")
+    try:
+        parsed_stats["temp"] = float(stats.get("Temperature (Sensor memory) (C)"))  # type: ignore
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU temperature as float")
+    try:
+        parsed_stats["powerWatts"] = float(
+            stats.get("Average Graphics Package Power (W)")  # type: ignore
+        )
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU power as float")
+    try:
+        parsed_stats["powerPercent"] = (
+            float(stats.get("Average Graphics Package Power (W)"))  # type: ignore
+            / float(stats.get("Max Graphics Package Power (W)"))  # type: ignore
+            * 100
+        )
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU average/max power as float")
+
+    return parsed_stats
+
+
 class GPUAMDStats:
     """Stats for AMD GPU devices."""
 
@@ -58,40 +92,6 @@ class GPUAMDStats:
     def __init__(self) -> None:
         self.samples = deque()
 
-    @staticmethod
-    def parse_stats(stats: Dict[str, str]) -> _Stats:
-        """Parse stats from rocm-smi output."""
-        parsed_stats: _Stats = {}
-
-        try:
-            parsed_stats["gpu"] = float(stats.get("GPU use (%)"))  # type: ignore
-        except (TypeError, ValueError):
-            logger.warning("Could not parse GPU usage as float")
-        try:
-            parsed_stats["memoryAllocated"] = float(stats.get("GPU memory use (%)"))  # type: ignore
-        except (TypeError, ValueError):
-            logger.warning("Could not parse GPU memory allocation as float")
-        try:
-            parsed_stats["temp"] = float(stats.get("Temperature (Sensor memory) (C)"))  # type: ignore
-        except (TypeError, ValueError):
-            logger.warning("Could not parse GPU temperature as float")
-        try:
-            parsed_stats["powerWatts"] = float(
-                stats.get("Average Graphics Package Power (W)")  # type: ignore
-            )
-        except (TypeError, ValueError):
-            logger.warning("Could not parse GPU power as float")
-        try:
-            parsed_stats["powerPercent"] = (
-                float(stats.get("Average Graphics Package Power (W)"))  # type: ignore
-                / float(stats.get("Max Graphics Package Power (W)"))  # type: ignore
-                * 100
-            )
-        except (TypeError, ValueError):
-            logger.warning("Could not parse GPU average/max power as float")
-
-        return parsed_stats
-
     def sample(self) -> None:
         try:
             raw_stats = get_rocm_smi_stats()
@@ -103,7 +103,7 @@ class GPUAMDStats:
 
             for card_key in card_keys:
                 card_stats = raw_stats[card_key]
-                stats = self.parse_stats(card_stats)
+                stats = parse_stats(card_stats)
                 if stats:
                     cards.append(stats)
 
@@ -183,7 +183,7 @@ class GPUAMD:
 
         can_read_rocm_smi = False
         try:
-            if get_rocm_smi_stats():
+            if parse_stats(get_rocm_smi_stats()):
                 can_read_rocm_smi = True
         except Exception:
             pass