wandb 0.17.5__py3-none-any.whl → 0.17.7__py3-none-any.whl

wandb/sdk/interface/interface.py

@@ -8,11 +8,14 @@ InterfaceRelay: Responses are routed to a relay queue (not matching uuids)
 
 """
 
+import gzip
 import logging
 import os
 import sys
 import time
 from abc import abstractmethod
+from pathlib import Path
+from secrets import token_hex
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -47,20 +50,24 @@ from ..lib.mailbox import MailboxHandle
 from . import summary_record as sr
 from .message_future import MessageFuture
 
+MANIFEST_FILE_SIZE_THRESHOLD = 100_000
+
 GlobStr = NewType("GlobStr", str)
 
-if TYPE_CHECKING:
-    from ..wandb_run import Run
+if sys.version_info >= (3, 8):
+    from typing import Literal, TypedDict
+else:
+    from typing_extensions import Literal, TypedDict
 
-    if sys.version_info >= (3, 8):
-        from typing import Literal, TypedDict
-    else:
-        from typing_extensions import Literal, TypedDict
+PolicyName = Literal["now", "live", "end"]
 
-    PolicyName = Literal["now", "live", "end"]
 
-    class FilesDict(TypedDict):
-        files: Iterable[Tuple[GlobStr, PolicyName]]
+class FilesDict(TypedDict):
+    files: Iterable[Tuple[GlobStr, PolicyName]]
+
+
+if TYPE_CHECKING:
+    from ..wandb_run import Run
 
 
 logger = logging.getLogger("wandb")
@@ -107,15 +114,14 @@ class InterfaceBase:
     def _publish_header(self, header: pb.HeaderRecord) -> None:
         raise NotImplementedError
 
-    def communicate_status(self) -> Optional[pb.StatusResponse]:
-        status = pb.StatusRequest()
-        resp = self._communicate_status(status)
-        return resp
+    def deliver_status(self) -> MailboxHandle:
+        return self._deliver_status(pb.StatusRequest())
 
     @abstractmethod
-    def _communicate_status(
-        self, status: pb.StatusRequest
-    ) -> Optional[pb.StatusResponse]:
+    def _deliver_status(
+        self,
+        status: pb.StatusRequest,
+    ) -> MailboxHandle:
         raise NotImplementedError
 
     def _make_config(
@@ -334,6 +340,12 @@ class InterfaceBase:
         proto_manifest.version = artifact_manifest.version()
         proto_manifest.storage_policy = artifact_manifest.storage_policy.name()
 
+        # Very large manifests need to be written to file to avoid protobuf size limits.
+        if len(artifact_manifest) > MANIFEST_FILE_SIZE_THRESHOLD:
+            path = self._write_artifact_manifest_file(artifact_manifest)
+            proto_manifest.manifest_file_path = path
+            return proto_manifest
+
         for k, v in artifact_manifest.storage_policy.config().items() or {}.items():
             cfg = proto_manifest.storage_policy_config.add()
             cfg.key = k
@@ -358,6 +370,18 @@ class InterfaceBase:
                 proto_extra.value_json = json.dumps(v)
         return proto_manifest
 
+    def _write_artifact_manifest_file(self, manifest: ArtifactManifest) -> str:
+        manifest_dir = Path(get_staging_dir()) / "artifact_manifests"
+        manifest_dir.mkdir(parents=True, exist_ok=True)
+        # It would be simpler to use `manifest.to_json()`, but that gets very slow for
+        # large manifests since it encodes the whole thing as a single JSON object.
+        filename = f"{time.time()}_{token_hex(8)}.manifest_contents.jl.gz"
+        manifest_file_path = manifest_dir / filename
+        with gzip.open(manifest_file_path, mode="wt", compresslevel=1) as f:
+            for entry in manifest.entries.values():
+                f.write(f"{json.dumps(entry.to_json())}\n")
+        return str(manifest_file_path)
+
     def deliver_link_artifact(
         self,
         run: "Run",

wandb/sdk/interface/interface_shared.py

@@ -299,7 +299,7 @@ class InterfaceShared(InterfaceBase):
         raise NotImplementedError
 
     def _communicate(
-        self, rec: pb.Record, timeout: Optional[int] = 5, local: Optional[bool] = None
+        self, rec: pb.Record, timeout: Optional[int] = 30, local: Optional[bool] = None
     ) -> Optional[pb.Result]:
         return self._communicate_async(rec, local=local).get(timeout=timeout)
 
@@ -421,15 +421,12 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(alert=proto_alert)
         self._publish(rec)
 
-    def _communicate_status(
-        self, status: pb.StatusRequest
-    ) -> Optional[pb.StatusResponse]:
+    def _deliver_status(
+        self,
+        status: pb.StatusRequest,
+    ) -> MailboxHandle:
         req = self._make_request(status=status)
-        resp = self._communicate(req, local=True)
-        if resp is None:
-            return None
-        assert resp.response.status_response
-        return resp.response.status_response
+        return self._deliver_record(req)
 
     def _publish_exit(self, exit_data: pb.RunExitRecord) -> None:
         rec = self._make_record(exit=exit_data)

wandb/sdk/internal/datastore.py

@@ -69,7 +69,7 @@ class DataStore:
 
     def __init__(self) -> None:
         self._opened_for_scan = False
-        self._fp: Optional["IO[Any]"] = None
+        self._fp: Optional[IO[Any]] = None
         self._index = 0
         self._flush_offset = 0
         self._size_bytes = 0

wandb/sdk/internal/handler.py

@@ -745,8 +745,6 @@ class HandleManager:
         self._respond_result(result)
 
     def handle_request_status(self, record: Record) -> None:
-        # TODO(mempressure): do something better?
-        assert record.control.req_resp
         result = proto_util._result_from_record(record)
         self._respond_result(result)
 

wandb/sdk/internal/internal.py

@@ -62,7 +62,7 @@ def wandb_internal(
 
     """
     # mark this process as internal
-    wandb._set_internal_process()
+    wandb._set_internal_process()  # type: ignore
     _setup_tracelog()
     started = time.time()
 

wandb/sdk/internal/job_builder.py

@@ -423,15 +423,18 @@ class JobBuilder:
         api: Api,
         build_context: Optional[str] = None,
         dockerfile: Optional[str] = None,
+        base_image: Optional[str] = None,
     ) -> Optional[Artifact]:
         """Build a job artifact from the current run.
 
         Arguments:
+            api (Api): The API object to use to create the job artifact.
             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.
+            base_image (Optional[str]): The base image used to run the job code.
 
         Returns:
             Optional[Artifact]: The job artifact if it was successfully built,
@@ -467,8 +470,6 @@ 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
@@ -520,6 +521,8 @@ class JobBuilder:
             source["build_context"] = build_context  # type: ignore[typeddict-item]
         if dockerfile:
             source["dockerfile"] = dockerfile  # type: ignore[typeddict-item]
+        if base_image:
+            source["base_image"] = base_image  # 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

wandb/sdk/internal/sender.py

@@ -1,5 +1,7 @@
 """sender."""
 
+import contextlib
+import gzip
 import json
 import logging
 import os
@@ -66,6 +68,7 @@ else:
 if TYPE_CHECKING:
     from wandb.proto.wandb_internal_pb2 import (
         ArtifactManifest,
+        ArtifactManifestEntry,
         ArtifactRecord,
         HttpResponse,
         LocalInfo,
@@ -105,22 +108,18 @@ def _framework_priority() -> Generator[Tuple[str, str], None, None]:
 
 def _manifest_json_from_proto(manifest: "ArtifactManifest") -> Dict:
     if manifest.version == 1:
-        contents = {
-            content.path: {
-                "digest": content.digest,
-                "birthArtifactID": content.birth_artifact_id
-                if content.birth_artifact_id
-                else None,
-                "ref": content.ref if content.ref else None,
-                "size": content.size if content.size is not None else None,
-                "local_path": content.local_path if content.local_path else None,
-                "skip_cache": content.skip_cache,
-                "extra": {
-                    extra.key: json.loads(extra.value_json) for extra in content.extra
-                },
+        if manifest.manifest_file_path:
+            contents = {}
+            with gzip.open(manifest.manifest_file_path, "rt") as f:
+                for line in f:
+                    entry_json = json.loads(line)
+                    path = entry_json.pop("path")
+                    contents[path] = entry_json
+        else:
+            contents = {
+                content.path: _manifest_entry_from_proto(content)
+                for content in manifest.contents
             }
-            for content in manifest.contents
-        }
     else:
         raise ValueError(f"unknown artifact manifest version: {manifest.version}")
 
@@ -135,6 +134,19 @@ def _manifest_json_from_proto(manifest: "ArtifactManifest") -> Dict:
     }
 
 
+def _manifest_entry_from_proto(entry: "ArtifactManifestEntry") -> Dict:
+    birth_artifact_id = entry.birth_artifact_id if entry.birth_artifact_id else None
+    return {
+        "digest": entry.digest,
+        "birthArtifactID": birth_artifact_id,
+        "ref": entry.ref if entry.ref else None,
+        "size": entry.size if entry.size is not None else None,
+        "local_path": entry.local_path if entry.local_path else None,
+        "skip_cache": entry.skip_cache,
+        "extra": {extra.key: json.loads(extra.value_json) for extra in entry.extra},
+    }
+
+
 class ResumeState:
     resumed: bool
     step: int
@@ -1586,6 +1598,10 @@ class SendManager:
         )
 
         self._job_builder._handle_server_artifact(res, artifact)
+
+        if artifact.manifest.manifest_file_path:
+            with contextlib.suppress(FileNotFoundError):
+                os.remove(artifact.manifest.manifest_file_path)
         return res
 
     def send_alert(self, record: "Record") -> None:

wandb/sdk/internal/tb_watcher.py

@@ -123,7 +123,7 @@ class TBWatcher:
         self._force = force
         # TODO(jhr): do we need locking in this queue?
         self._watcher_queue = queue.PriorityQueue()
-        wandb.tensorboard.reset_state()
+        wandb.tensorboard.reset_state()  # type: ignore
 
     def _calculate_namespace(self, logdir: str, rootdir: str) -> Optional[str]:
         namespace: Optional[str]
@@ -430,7 +430,7 @@ class TBEventConsumer:
     def _handle_event(
         self, event: "ProtoEvent", history: Optional["TBHistory"] = None
     ) -> None:
-        wandb.tensorboard._log(
+        wandb.tensorboard._log(  # type: ignore
             event.event,
             step=event.event.step,
             namespace=event.namespace,

wandb/sdk/internal/update.py

@@ -10,7 +10,7 @@ def _find_available(
 ) -> Optional[Tuple[str, bool, bool, bool, Optional[str]]]:
     from wandb.util import parse_version
 
-    pypi_url = f"https://pypi.org/pypi/{wandb._wandb_module}/json"
+    pypi_url = "https://pypi.org/pypi/wandb/json"
 
     yanked_dict = {}
     try:
@@ -78,7 +78,7 @@ def check_available(current_version: str) -> Optional[Dict[str, Optional[str]]]:
     if not package_info:
         return None
 
-    wandb_module_name = wandb._wandb_module
+    wandb_module_name = "wandb"
 
     latest_version, pip_prerelease, deleted, yanked, yanked_reason = package_info
     upgrade_message = (

wandb/sdk/launch/_launch.py

@@ -211,7 +211,9 @@ async def _launch(
     launch_project = LaunchProject.from_spec(launch_spec, api)
     launch_project.fetch_and_validate_project()
     entrypoint = launch_project.get_job_entry_point()
-    image_uri = launch_project.docker_image  # Either set by user or None.
+    image_uri = (
+        launch_project.docker_image or launch_project.job_base_image
+    )  # Either set by user or None.
 
     # construct runner config.
     runner_config: Dict[str, Any] = {}
@@ -224,7 +226,7 @@ async def _launch(
         await environment.verify()
     registry = loader.registry_from_config(registry_config, environment)
     builder = loader.builder_from_config(build_config, environment, registry)
-    if not launch_project.docker_image:
+    if not (launch_project.docker_image or launch_project.job_base_image):
         assert entrypoint
         image_uri = await builder.build_image(launch_project, entrypoint, None)
     backend = loader.runner_from_config(

wandb/sdk/launch/_project_spec.py

@@ -7,6 +7,7 @@ import enum
 import json
 import logging
 import os
+import shutil
 import tempfile
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
@@ -112,6 +113,9 @@ class LaunchProject:
         self.sweep_id = sweep_id
         self.author = launch_spec.get("author")
         self.python_version: Optional[str] = launch_spec.get("python_version")
+        self._job_dockerfile: Optional[str] = None
+        self._job_build_context: Optional[str] = None
+        self._job_base_image: Optional[str] = None
         self.accelerator_base_image: Optional[str] = resource_args_build.get(
             "accelerator", {}
         ).get("base_image") or resource_args_build.get("cuda", {}).get("base_image")
@@ -131,8 +135,6 @@ class LaunchProject:
         self._queue_name: Optional[str] = None
         self._queue_entity: Optional[str] = None
         self._run_queue_item_id: Optional[str] = None
-        self._job_dockerfile: Optional[str] = None
-        self._job_build_context: Optional[str] = None
 
     def init_source(self) -> None:
         if self.docker_image is not None:
@@ -146,6 +148,21 @@ class LaunchProject:
             self.project_dir = os.getcwd()
             self._entry_point = self.override_entrypoint
 
+    def change_project_dir(self, new_dir: str) -> None:
+        """Change the project directory to a new directory."""
+        # Copy the contents of the old project dir to the new project dir.
+        old_dir = self.project_dir
+        if old_dir is not None:
+            shutil.copytree(
+                old_dir,
+                new_dir,
+                symlinks=True,
+                dirs_exist_ok=True,
+                ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc", ".git"),
+            )
+            shutil.rmtree(old_dir)
+        self.project_dir = new_dir
+
     def init_git(self, git_info: Dict[str, str]) -> None:
         self.git_version = git_info.get("version")
         self.git_repo = git_info.get("repo")
@@ -212,14 +229,23 @@ class LaunchProject:
     def job_build_context(self) -> Optional[str]:
         return self._job_build_context
 
+    @property
+    def job_base_image(self) -> Optional[str]:
+        return self._job_base_image
+
     def set_job_dockerfile(self, dockerfile: str) -> None:
         self._job_dockerfile = dockerfile
 
     def set_job_build_context(self, build_context: str) -> None:
         self._job_build_context = build_context
 
+    def set_job_base_image(self, base_image: str) -> None:
+        self._job_base_image = base_image
+
     @property
     def image_name(self) -> str:
+        if self.job_base_image is not None:
+            return self.job_base_image
         if self.docker_image is not None:
             return self.docker_image
         elif self.uri is not None:
@@ -299,10 +325,8 @@ class LaunchProject:
 
     def build_required(self) -> bool:
         """Checks the source to see if a build is required."""
-        # since the image tag for images built from jobs
-        # is based on the job version index, which is immutable
-        # we don't need to build the image for a job if that tag
-        # already exists
+        if self.job_base_image is not None:
+            return False
         if self.source != LaunchSource.JOB:
             return True
         return False
@@ -316,7 +340,9 @@ class LaunchProject:
         Returns:
             Optional[str]: The Docker image or None if not specified.
         """
-        return self._docker_image
+        if self._docker_image:
+            return self._docker_image
+        return None
 
     @docker_image.setter
     def docker_image(self, value: str) -> None:
@@ -336,7 +362,7 @@ class LaunchProject:
         # assuming project only has 1 entry point, pull that out
         # tmp fn until we figure out if we want to support multiple entry points or not
         if not self._entry_point:
-            if not self.docker_image:
+            if not self.docker_image and not self.job_base_image:
                 raise LaunchError(
                     "Project must have at least one entry point unless docker image is specified."
                 )

wandb/sdk/launch/agent/agent.py

@@ -717,7 +717,7 @@ class LaunchAgent:
         _, build_config, registry_config = construct_agent_configs(
             default_config, override_build_config
         )
-        image_uri = project.docker_image
+        image_uri = project.docker_image or project.job_base_image
         entrypoint = project.get_job_entry_point()
         environment = loader.environment_from_config(
             default_config.get("environment", {})
@@ -727,7 +727,11 @@ class LaunchAgent:
         backend = loader.runner_from_config(
             resource, api, backend_config, environment, registry
         )
-        if not (project.docker_image or isinstance(backend, LocalProcessRunner)):
+        if not (
+            project.docker_image
+            or project.job_base_image
+            or isinstance(backend, LocalProcessRunner)
+        ):
             assert entrypoint is not None
             image_uri = await builder.build_image(project, entrypoint, job_tracker)
 

wandb/sdk/launch/agent/run_queue_item_file_saver.py

@@ -2,7 +2,7 @@
 
 import os
 import sys
-from typing import List, Optional, Union
+from typing import List, Optional
 
 import wandb
 
@@ -17,9 +17,7 @@ FileSubtypes = Literal["warning", "error"]
 class RunQueueItemFileSaver:
     def __init__(
         self,
-        agent_run: Optional[
-            Union["wandb.sdk.wandb_run.Run", "wandb.sdk.lib.RunDisabled"]
-        ],
+        agent_run: Optional["wandb.sdk.wandb_run.Run"],
         run_queue_item_id: str,
     ):
         self.run_queue_item_id = run_queue_item_id

wandb/sdk/launch/builder/build.py

@@ -201,7 +201,7 @@ def get_requirements_section(
     # If there is a requirements.txt at root of build context, use that.
     if (base_path / "src" / "requirements.txt").exists():
         requirements_files += ["src/requirements.txt"]
-        deps_install_line = "pip install -r requirements.txt"
+        deps_install_line = "pip install uv && uv pip install -r requirements.txt"
         with open(base_path / "src" / "requirements.txt") as f:
             requirements = f.readlines()
         if not any(["wandb" in r for r in requirements]):
@@ -237,7 +237,9 @@ def get_requirements_section(
                 with open(base_path / "src" / "requirements.txt", "w") as f:
                     f.write("\n".join(project_deps))
                 requirements_files += ["src/requirements.txt"]
-                deps_install_line = "pip install -r requirements.txt"
+                deps_install_line = (
+                    "pip install uv && uv pip install -r requirements.txt"
+                )
                 return PIP_TEMPLATE.format(
                     buildx_optional_prefix=prefix,
                     requirements_files=" ".join(requirements_files),

wandb/sdk/launch/builder/kaniko_builder.py

@@ -63,6 +63,13 @@ else:
     NAMESPACE = "wandb"
 
 
+def get_pod_name_safe(job: client.V1Job):
+    try:
+        return job.spec.template.metadata.name
+    except AttributeError:
+        return None
+
+
 async def _wait_for_completion(
     batch_client: client.BatchV1Api, job_name: str, deadline_secs: Optional[int] = None
 ) -> bool:
@@ -319,17 +326,18 @@ class KanikoBuilder(AbstractBuilder):
                 await self._create_docker_ecr_config_map(
                     build_job_name, core_v1, repo_uri
                 )
-            await batch_v1.create_namespaced_job(NAMESPACE, build_job)
-
+            k8s_job = await batch_v1.create_namespaced_job(NAMESPACE, build_job)
             # wait for double the job deadline since it might take time to schedule
             if not await _wait_for_completion(
                 batch_v1, build_job_name, 3 * _DEFAULT_BUILD_TIMEOUT_SECS
             ):
                 if job_tracker:
                     job_tracker.set_err_stage("build")
-                raise Exception(
-                    f"Failed to build image in kaniko for job {run_id}. View logs with `kubectl logs -n {NAMESPACE} {build_job_name}`."
-                )
+                msg = f"Failed to build image in kaniko for job {run_id}."
+                pod_name = get_pod_name_safe(k8s_job)
+                if pod_name:
+                    msg += f" View logs with `kubectl logs -n {NAMESPACE} {pod_name}`."
+                raise Exception(msg)
             try:
                 pods_from_job = await core_v1.list_namespaced_pod(
                     namespace=NAMESPACE, label_selector=f"job-name={build_job_name}"

wandb/sdk/launch/builder/templates/_wandb_bootstrap.py

@@ -39,12 +39,13 @@ def install_deps(
         deps (str[], None): The dependencies that failed to install
     """
     try:
+        subprocess.check_output(["pip", "install", "uv"], stderr=subprocess.STDOUT)
         # Include only uri if @ is present
         clean_deps = [d.split("@")[-1].strip() if "@" in d else d for d in deps]
         index_args = ["--extra-index-url", extra_index] if extra_index else []
         print("installing {}...".format(", ".join(clean_deps)))
         opts = opts or []
-        args = ["pip", "install"] + opts + clean_deps + index_args
+        args = ["uv", "pip", "install"] + opts + clean_deps + index_args
         sys.stdout.flush()
         subprocess.check_output(args, stderr=subprocess.STDOUT)
         return failed

wandb/sdk/launch/create_job.py

@@ -114,6 +114,7 @@ def _create_job(
     git_hash: Optional[str] = None,
     build_context: Optional[str] = None,
     dockerfile: Optional[str] = None,
+    base_image: Optional[str] = None,
 ) -> Tuple[Optional[Artifact], str, List[str]]:
     wandb.termlog(f"Creating launch job of type: {job_type}...")
 
@@ -188,6 +189,7 @@ def _create_job(
         api.api,
         dockerfile=dockerfile,
         build_context=build_context,
+        base_image=base_image,
     )
     if not artifact:
         wandb.termerror("JobBuilder failed to build a job")