wandb 0.15.4__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/interface/interface.py

@@ -20,8 +20,6 @@ from typing import TYPE_CHECKING, Any, Iterable, NewType, Optional, Tuple, Union
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_telemetry_pb2 as tpb
 from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
-from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
 from wandb.util import (
     WandBJSONEncoderOld,
     get_h5_typename,
@@ -40,6 +38,8 @@ from .message_future import MessageFuture
 GlobStr = NewType("GlobStr", str)
 
 if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+
     from ..wandb_run import Run
 
     if sys.version_info >= (3, 8):
@@ -373,7 +373,7 @@ class InterfaceBase:
     def _publish_files(self, files: pb.FilesRecord) -> None:
         raise NotImplementedError
 
-    def _make_artifact(self, artifact: LocalArtifact) -> pb.ArtifactRecord:
+    def _make_artifact(self, artifact: "Artifact") -> pb.ArtifactRecord:
         proto_artifact = pb.ArtifactRecord()
         proto_artifact.type = artifact.type
         proto_artifact.name = artifact.name
@@ -425,14 +425,14 @@ class InterfaceBase:
     def publish_link_artifact(
         self,
         run: "Run",
-        artifact: Union[LocalArtifact, PublicArtifact],
+        artifact: "Artifact",
         portfolio_name: str,
         aliases: Iterable[str],
         entity: Optional[str] = None,
         project: Optional[str] = None,
     ) -> None:
         link_artifact = pb.LinkArtifactRecord()
-        if isinstance(artifact, LocalArtifact):
+        if artifact.is_draft():
             link_artifact.client_id = artifact._client_id
         else:
             link_artifact.server_id = artifact.id if artifact.id else ""
@@ -449,10 +449,8 @@ class InterfaceBase:
 
     def publish_use_artifact(
         self,
-        artifact: LocalArtifact,
+        artifact: "Artifact",
     ) -> None:
-        # use_artifact is either a public.Artifact or a wandb.Artifact that has been
-        # waited on and has an id
         assert artifact.id is not None, "Artifact must have an id"
         use_artifact = pb.UseArtifactRecord(
             id=artifact.id, type=artifact.type, name=artifact.name
@@ -467,7 +465,7 @@ class InterfaceBase:
     def communicate_artifact(
         self,
         run: "Run",
-        artifact: LocalArtifact,
+        artifact: "Artifact",
         aliases: Iterable[str],
         history_step: Optional[int] = None,
         is_user_created: bool = False,
@@ -517,7 +515,7 @@ class InterfaceBase:
     def publish_artifact(
         self,
         run: "Run",
-        artifact: LocalArtifact,
+        artifact: "Artifact",
         aliases: Iterable[str],
         is_user_created: bool = False,
         use_after_commit: bool = False,

wandb/sdk/internal/file_stream.py

@@ -653,9 +653,8 @@ def request_with_retry(
                 e.response is not None and e.response.status_code == 429
             ):
                 err_str = (
-                    "Filestream rate limit exceeded, retrying in {} seconds".format(
-                        delay
-                    )
+                    "Filestream rate limit exceeded, "
+                    f"retrying in {delay:.1f} seconds. "
                 )
                 if retry_callback:
                     retry_callback(e.response.status_code, err_str)

wandb/sdk/internal/internal_api.py

@@ -287,6 +287,7 @@ class Api:
         self.server_use_artifact_input_info: Optional[List[str]] = None
         self._max_cli_version: Optional[str] = None
         self._server_settings_type: Optional[List[str]] = None
+        self.fail_run_queue_item_input_info: Optional[List[str]] = None
 
     def gql(self, *args: Any, **kwargs: Any) -> Any:
         ret = self._retry_gql(
@@ -586,13 +587,103 @@ class Api:
         return "failRunQueueItem" in mutations
 
     @normalize_exceptions
-    def fail_run_queue_item(self, run_queue_item_id: str) -> bool:
+    def fail_run_queue_item_fields_introspection(self) -> List:
+        if self.fail_run_queue_item_input_info:
+            return self.fail_run_queue_item_input_info
+        query_string = """
+           query ProbeServerFailRunQueueItemInput {
+                FailRunQueueItemInputInfoType: __type(name:"FailRunQueueItemInput") {
+                    inputFields{
+                        name
+                    }
+                }
+            }
+        """
+
+        query = gql(query_string)
+        res = self.gql(query)
+
+        self.fail_run_queue_item_input_info = [
+            field.get("name", "")
+            for field in res.get("FailRunQueueItemInputInfoType", {}).get(
+                "inputFields", [{}]
+            )
+        ]
+        return self.fail_run_queue_item_input_info
+
+    @normalize_exceptions
+    def fail_run_queue_item(
+        self,
+        run_queue_item_id: str,
+        message: str,
+        stage: str,
+        file_paths: Optional[List[str]] = None,
+    ) -> bool:
+        if not self.fail_run_queue_item_introspection():
+            return False
+        variable_values: Dict[str, Union[str, Optional[List[str]]]] = {
+            "runQueueItemId": run_queue_item_id,
+        }
+        if "message" in self.fail_run_queue_item_fields_introspection():
+            variable_values.update({"message": message, "stage": stage})
+            if file_paths is not None:
+                variable_values["filePaths"] = file_paths
+            mutation_string = """
+            mutation failRunQueueItem($runQueueItemId: ID!, $message: String!, $stage: String!, $filePaths: [String!]) {
+                failRunQueueItem(
+                    input: {
+                        runQueueItemId: $runQueueItemId
+                        message: $message
+                        stage: $stage
+                        filePaths: $filePaths
+                    }
+                ) {
+                    success
+                }
+            }
+            """
+        else:
+            mutation_string = """
+            mutation failRunQueueItem($runQueueItemId: ID!) {
+                failRunQueueItem(
+                    input: {
+                        runQueueItemId: $runQueueItemId
+                    }
+                ) {
+                    success
+                }
+            }
+            """
+
+        mutation = gql(mutation_string)
+        response = self.gql(mutation, variable_values=variable_values)
+        result: bool = response["failRunQueueItem"]["success"]
+        return result
+
+    @normalize_exceptions
+    def update_run_queue_item_warning_introspection(self) -> bool:
+        _, _, mutations = self.server_info_introspection()
+        return "updateRunQueueItemWarning" in mutations
+
+    @normalize_exceptions
+    def update_run_queue_item_warning(
+        self,
+        run_queue_item_id: str,
+        message: str,
+        stage: str,
+        file_paths: Optional[List[str]] = None,
+    ) -> bool:
+        if not self.update_run_queue_item_warning_introspection():
+            return False
         mutation = gql(
             """
-        mutation failRunQueueItem($runQueueItemId: ID!) {
-            failRunQueueItem(
+        mutation updateRunQueueItemWarning($runQueueItemId: ID!, $message: String!, $stage: String!, $filePaths: [String!]) {
+            updateRunQueueItemWarning(
                 input: {
                     runQueueItemId: $runQueueItemId
+                    message: $message
+                    stage: $stage
+                    filePaths: $filePaths
                 }
             ) {
                 success
@@ -604,9 +695,12 @@ class Api:
             mutation,
             variable_values={
                 "runQueueItemId": run_queue_item_id,
+                "message": message,
+                "stage": stage,
+                "filePaths": file_paths,
             },
         )
-        result: bool = response["failRunQueueItem"]["success"]
+        result: bool = response["updateRunQueueItemWarning"]["success"]
         return result
 
     @normalize_exceptions
@@ -617,6 +711,7 @@ class Api:
             viewer {
                 id
                 entity
+                username
                 flags
                 teams {
                     edges {

wandb/sdk/internal/job_builder.py

@@ -5,7 +5,7 @@ import os
 import sys
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
-from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.data_types._dtypes import TypeRegistry
 from wandb.sdk.lib.filenames import DIFF_FNAME, METADATA_FNAME, REQUIREMENTS_FNAME
 from wandb.util import make_artifact_name_safe
@@ -62,7 +62,7 @@ class ArtifactInfoForJob(TypedDict):
     name: str
 
 
-class JobArtifact(LocalArtifact):
+class JobArtifact(Artifact):
     def __init__(self, name: str, *args: Any, **kwargs: Any):
         super().__init__(name, "placeholder", *args, **kwargs)
         self._type = JOB_ARTIFACT_TYPE  # Get around type restriction.
@@ -76,6 +76,7 @@ class JobBuilder:
     _summary: Optional[Dict[str, Any]]
     _logged_code_artifact: Optional[ArtifactInfoForJob]
     _disable: bool
+    _aliases: List[str]
 
     def __init__(self, settings: SettingsStatic):
         self._settings = settings
@@ -85,6 +86,7 @@ class JobBuilder:
         self._summary = None
         self._logged_code_artifact = None
         self._disable = settings.disable_job_creation
+        self._aliases = []
         self._source_type: Optional[
             Literal["repo", "artifact", "image"]
         ] = settings.get("job_source")
@@ -116,7 +118,7 @@ class JobBuilder:
 
     def _build_repo_job(
         self, metadata: Dict[str, Any], program_relpath: str, root: Optional[str]
-    ) -> Tuple[Optional[LocalArtifact], Optional[GitSourceDict]]:
+    ) -> Tuple[Optional[Artifact], Optional[GitSourceDict]]:
         git_info: Dict[str, str] = metadata.get("git", {})
         remote = git_info.get("remote")
         commit = git_info.get("commit")
@@ -177,7 +179,7 @@ class JobBuilder:
 
     def _build_artifact_job(
         self, metadata: Dict[str, Any], program_relpath: str
-    ) -> Tuple[Optional[LocalArtifact], Optional[ArtifactSourceDict]]:
+    ) -> Tuple[Optional[Artifact], Optional[ArtifactSourceDict]]:
         assert isinstance(self._logged_code_artifact, dict)
         # TODO: should we just always exit early if the path doesn't exist?
         if self._is_notebook_run() and not self._is_colab_run():
@@ -212,10 +214,16 @@ class JobBuilder:
 
     def _build_image_job(
         self, metadata: Dict[str, Any]
-    ) -> Tuple[LocalArtifact, ImageSourceDict]:
+    ) -> Tuple[Artifact, ImageSourceDict]:
         image_name = metadata.get("docker")
         assert isinstance(image_name, str)
-        name = make_artifact_name_safe(f"job-{image_name}")
+
+        raw_image_name = image_name
+        if ":" in image_name:
+            raw_image_name, tag = image_name.split(":")
+            self._aliases += [tag]
+
+        name = make_artifact_name_safe(f"job-{raw_image_name}")
         artifact = JobArtifact(name)
         source: ImageSourceDict = {
             "image": image_name,
@@ -228,7 +236,7 @@ class JobBuilder:
     def _is_colab_run(self) -> bool:
         return hasattr(self._settings, "_colab") and bool(self._settings._colab)
 
-    def build(self) -> Optional[LocalArtifact]:
+    def build(self) -> Optional[Artifact]:
         _logger.info("Attempting to build job artifact")
         if not os.path.exists(
             os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME)

wandb/sdk/internal/sender.py

@@ -1608,6 +1608,10 @@ class SendManager:
             # TODO: this should be removed when the latest tag is handled
             # by the backend (WB-12116)
             proto_artifact.aliases.append("latest")
+            # add docker image tag
+            for alias in self._job_builder._aliases:
+                proto_artifact.aliases.append(alias)
+
             proto_artifact.user_created = True
             proto_artifact.use_after_commit = True
             proto_artifact.finalize = True

wandb/sdk/internal/settings_static.py

@@ -8,6 +8,7 @@ class SettingsStatic:
     # TODO(jhr): figure out how to share type defs with sdk/wandb_settings.py
     _offline: Optional[bool]
     _sync: bool
+    _disable_setproctitle: bool
     _disable_stats: Optional[bool]
     _disable_meta: Optional[bool]
     _flow_control: bool

wandb/sdk/launch/_project_spec.py

@@ -7,19 +7,21 @@ import json
 import logging
 import os
 import tempfile
-from typing import Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import wandb
 import wandb.docker as docker
 from wandb.apis.internal import Api
 from wandb.errors import CommError
-from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
 from wandb.sdk.launch import utils
 from wandb.sdk.lib.runid import generate_id
 
 from .errors import LaunchError
 from .utils import LOG_PREFIX, recursive_macro_sub
 
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+
 _logger = logging.getLogger(__name__)
 
 DEFAULT_LAUNCH_METADATA_PATH = "launch_metadata.json"
@@ -69,7 +71,7 @@ class LaunchProject:
         self.job = job
         if job is not None:
             wandb.termlog(f"{LOG_PREFIX}Launching job: {job}")
-        self._job_artifact: Optional[PublicArtifact] = None
+        self._job_artifact: Optional["Artifact"] = None
         self.api = api
         self.launch_spec = launch_spec
         self.target_entity = target_entity
@@ -80,12 +82,12 @@ class LaunchProject:
         # runner, so we need to pop the builder key out
         resource_args_build = resource_args.get(resource, {}).pop("builder", {})
         self.resource = resource
-        self.resource_args = resource_args
+        self.resource_args = resource_args.copy()
         self.sweep_id = sweep_id
         self.python_version: Optional[str] = launch_spec.get("python_version")
-        self.cuda_base_image: Optional[str] = resource_args_build.get("cuda", {}).get(
-            "base_image"
-        )
+        self.accelerator_base_image: Optional[str] = resource_args_build.get(
+            "accelerator", {}
+        ).get("base_image") or resource_args_build.get("cuda", {}).get("base_image")
         self._base_image: Optional[str] = launch_spec.get("base_image")
         self.docker_image: Optional[str] = docker_config.get(
             "docker_image"