wandb 0.17.8rc1__py3-none-any.whl → 0.18.0__py3-none-any.whl

wandb/sdk/artifacts/storage_handlers/gcs_handler.py

@@ -19,6 +19,12 @@ if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact import Artifact
 
 
+class _GCSIsADirectoryError(Exception):
+    """Raised when we try to download a GCS folder."""
+
+    pass
+
+
 class GCSHandler(StorageHandler):
     _client: Optional["gcs_module.client.Client"]
 
@@ -70,6 +76,11 @@ class GCSHandler(StorageHandler):
         bucket, key, _ = self._parse_uri(manifest_entry.ref)
         version = manifest_entry.extra.get("versionID")
 
+        if self._is_dir(manifest_entry):
+            raise _GCSIsADirectoryError(
+                f"Unable to download GCS folder {manifest_entry.ref!r}, skipping"
+            )
+
         obj = None
         # First attempt to get the generation specified, this will return None if versioning is not enabled
         if version is not None:
@@ -135,6 +146,7 @@ class GCSHandler(StorageHandler):
         entries = [
             self._entry_from_obj(obj, path, name, prefix=key, multi=multi)
             for obj in objects
+            if not obj.name.endswith("/")
         ]
         if start_time is not None:
             termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
@@ -193,3 +205,22 @@ class GCSHandler(StorageHandler):
             size=obj.size,
             extra={"versionID": obj.generation},
         )
+
+    def _is_dir(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+    ) -> bool:
+        assert self._client is not None
+        assert manifest_entry.ref is not None
+        bucket, key, _ = self._parse_uri(manifest_entry.ref)
+        bucket_obj = self._client.bucket(bucket)
+        # A gcs bucket key should end with a forward slash on gcloud, but
+        # we save these refs without the forward slash in the manifest entry
+        # so we check the size and extension, make sure its not referring to
+        # an actual file with this reference, and that the ref with the slash
+        # exists on gcloud
+        return key.endswith("/") or (
+            not (manifest_entry.size or PurePosixPath(key).suffix)
+            and bucket_obj.get_blob(key) is None
+            and bucket_obj.get_blob(f"{key}/") is not None
+        )

wandb/sdk/interface/interface.py

@@ -519,6 +519,7 @@ class InterfaceBase:
         run: "Run",
         artifact: "Artifact",
         aliases: Iterable[str],
+        tags: Optional[Iterable[str]] = None,
         history_step: Optional[int] = None,
         is_user_created: bool = False,
         use_after_commit: bool = False,
@@ -532,8 +533,9 @@ class InterfaceBase:
         proto_artifact.user_created = is_user_created
         proto_artifact.use_after_commit = use_after_commit
         proto_artifact.finalize = finalize
-        for alias in aliases:
-            proto_artifact.aliases.append(alias)
+
+        proto_artifact.aliases.extend(aliases or [])
+        proto_artifact.tags.extend(tags or [])
 
         log_artifact = pb.LogArtifactRequest()
         log_artifact.artifact.CopyFrom(proto_artifact)
@@ -577,6 +579,7 @@ class InterfaceBase:
         run: "Run",
         artifact: "Artifact",
         aliases: Iterable[str],
+        tags: Optional[Iterable[str]] = None,
         is_user_created: bool = False,
         use_after_commit: bool = False,
         finalize: bool = True,
@@ -589,8 +592,8 @@ class InterfaceBase:
         proto_artifact.user_created = is_user_created
         proto_artifact.use_after_commit = use_after_commit
         proto_artifact.finalize = finalize
-        for alias in aliases:
-            proto_artifact.aliases.append(alias)
+        proto_artifact.aliases.extend(aliases or [])
+        proto_artifact.tags.extend(tags or [])
         self._publish_artifact(proto_artifact)
 
     @abstractmethod

wandb/sdk/internal/internal_api.py

@@ -3551,7 +3551,7 @@ class Api:
         _id: Optional[str] = response["createArtifactType"]["artifactType"]["id"]
         return _id
 
-    def server_artifact_introspection(self) -> List:
+    def server_artifact_introspection(self) -> List[str]:
         query_string = """
             query ProbeServerArtifact {
                 ArtifactInfoType: __type(name:"Artifact") {
@@ -3572,7 +3572,7 @@ class Api:
 
         return self.server_artifact_fields_info
 
-    def server_create_artifact_introspection(self) -> List:
+    def server_create_artifact_introspection(self) -> List[str]:
         query_string = """
             query ProbeServerCreateArtifactInput {
                 CreateArtifactInputInfoType: __type(name:"CreateArtifactInput") {
@@ -3627,6 +3627,10 @@ class Api:
             types += "$ttlDurationSeconds: Int64,"
             values += "ttlDurationSeconds: $ttlDurationSeconds,"
 
+        if "tags" in fields:
+            types += "$tags: [TagInput!],"
+            values += "tags: $tags,"
+
         query_template = """
             mutation CreateArtifact(
                 $artifactTypeName: String!,
@@ -3686,18 +3690,25 @@ class Api:
         metadata: Optional[Dict] = None,
         ttl_duration_seconds: Optional[int] = None,
         aliases: Optional[List[Dict[str, str]]] = None,
+        tags: Optional[List[Dict[str, str]]] = None,
         distributed_id: Optional[str] = None,
         is_user_created: Optional[bool] = False,
         history_step: Optional[int] = None,
     ) -> Tuple[Dict, Dict]:
         fields = self.server_create_artifact_introspection()
         artifact_fields = self.server_artifact_introspection()
-        if "ttlIsInherited" not in artifact_fields and ttl_duration_seconds:
+        if ("ttlIsInherited" not in artifact_fields) and ttl_duration_seconds:
             wandb.termwarn(
                 "Server not compatible with setting Artifact TTLs, please upgrade the server to use Artifact TTL"
             )
             # ttlDurationSeconds is only usable if ttlIsInherited is also present
             ttl_duration_seconds = None
+        if ("tags" not in artifact_fields) and tags:
+            wandb.termwarn(
+                "Server not compatible with Artifact tags. "
+                "To use Artifact tags, please upgrade the server to v0.85 or higher."
+            )
+
         query_template = self._get_create_artifact_mutation(
             fields, history_step, distributed_id
         )
@@ -3706,8 +3717,6 @@ class Api:
         project_name = project_name or self.settings("project")
         if not is_user_created:
             run_name = run_name or self.current_run_id
-        if aliases is None:
-            aliases = []
 
         mutation = gql(query_template)
         response = self.gql(
@@ -3722,7 +3731,8 @@ class Api:
                 "sequenceClientID": sequence_client_id,
                 "digest": digest,
                 "description": description,
-                "aliases": [alias for alias in aliases],
+                "aliases": list(aliases or []),
+                "tags": list(tags or []),
                 "metadata": json.dumps(util.make_safe_for_json(metadata))
                 if metadata
                 else None,

wandb/sdk/internal/sender.py

@@ -1589,6 +1589,7 @@ class SendManager:
             ttl_duration_seconds=artifact.ttl_duration_seconds or None,
             description=artifact.description or None,
             aliases=artifact.aliases,
+            tags=artifact.tags,
             use_after_commit=artifact.use_after_commit,
             distributed_id=artifact.distributed_id,
             finalize=artifact.finalize,

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

@@ -183,8 +183,17 @@ class GPUAMD:
 
         can_read_rocm_smi = False
         try:
-            if parse_stats(get_rocm_smi_stats()):
-                can_read_rocm_smi = True
+            # try to read stats from rocm-smi and parse them
+            raw_stats = get_rocm_smi_stats()
+            card_keys = [
+                key for key in sorted(raw_stats.keys()) if key.startswith("card")
+            ]
+
+            for card_key in card_keys:
+                card_stats = raw_stats[card_key]
+                parse_stats(card_stats)
+
+            can_read_rocm_smi = True
         except Exception:
             pass
 

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

@@ -224,7 +224,8 @@ class NeuronCoreStats:
                 for k, v in usage_breakdown["neuroncore_memory_usage"].items()
             }
 
-            # executed with torchrun? only keep the local_rank stats
+            # When the training script is executed with torchrun,
+            # we only want to keep the relevant LOCAL_RANK stats
             local_rank = int(os.environ.get("LOCAL_RANK", -1337))
             if local_rank >= 0:
                 neuroncore_utilization = {

wandb/sdk/internal/tb_watcher.py

@@ -14,7 +14,7 @@ import wandb
 from wandb import util
 from wandb.sdk.interface.interface import GlobStr
 from wandb.sdk.lib import filesystem
-from wandb.viz import CustomChart
+from wandb.sdk.lib.viz import CustomChart
 
 from . import run as internal_run
 

wandb/sdk/launch/inputs/internal.py

@@ -141,9 +141,9 @@ def _replace_refs_and_allofs(schema: dict, defs: Optional[dict]) -> dict:
     ret: Dict[str, Any] = {}
     if "$ref" in schema and defs:
         # Reference found, replace it with its definition
-        def_key = schema["$ref"].split("#/$defs/")[1]
+        def_key = schema.pop("$ref").split("#/$defs/")[1]
         # Also run recursive replacement in case a ref contains more refs
-        return _replace_refs_and_allofs(defs.pop(def_key), defs)
+        ret = _replace_refs_and_allofs(defs.pop(def_key), defs)
     for key, val in schema.items():
         if isinstance(val, dict):
             # Step into dicts recursively

wandb/sdk/lib/_settings_toposort_generated.py

@@ -62,7 +62,7 @@ _Setting = Literal[
     "_proxies",
     "_python",
     "_runqueue_item_id",
-    "_require_core",
+    "_require_legacy_service",
     "_save_requirements",
     "_service_transport",
     "_service_wait",
@@ -70,6 +70,7 @@ _Setting = Literal[
     "_start_datetime",
     "_start_time",
     "_stats_pid",
+    "_stats_sampling_interval",
     "_stats_sample_rate_seconds",
     "_stats_samples_to_average",
     "_stats_join_assets",
@@ -173,8 +174,6 @@ _Setting = Literal[
     "sync_dir",
     "sync_file",
     "sync_symlink_latest",
-    "system_sample",
-    "system_sample_seconds",
     "table_raise_on_max_row_limit_exceeded",
     "timespec",
     "tmp_dir",
@@ -186,6 +185,7 @@ SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
     "_service_wait",
     "_stats_sample_rate_seconds",
     "_stats_samples_to_average",
+    "_stats_sampling_interval",
     "anonymous",
     "api_key",
     "base_url",

wandb/sdk/service/server_sock.py

@@ -199,7 +199,6 @@ class SockAcceptThread(threading.Thread):
         self._clients = ClientDict()
 
     def run(self) -> None:
-        self._sock.listen(5)
         read_threads = []
 
         while not self._stopped.is_set():
@@ -254,6 +253,7 @@ class SocketServer:
 
     def start(self) -> None:
         self._bind()
+        self._sock.listen(5)
         self._thread = SockAcceptThread(sock=self._sock, mux=self._mux)
         self._thread.start()
         # Note: Uncomment to figure out what thread is not exiting properly

wandb/sdk/service/service.py

@@ -15,7 +15,11 @@ import time
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from wandb import _sentry, termlog
-from wandb.env import core_debug, core_error_reporting_enabled, is_require_core
+from wandb.env import (
+    core_debug,
+    core_error_reporting_enabled,
+    is_require_legacy_service,
+)
 from wandb.errors import Error, WandbCoreNotAvailableError
 from wandb.sdk.lib.wburls import wburls
 from wandb.util import get_core_path, get_module
@@ -164,7 +168,7 @@ class _Service:
 
             service_args = []
 
-            if is_require_core():
+            if not is_require_legacy_service():
                 try:
                     core_path = get_core_path()
                 except WandbCoreNotAvailableError as e:

wandb/sdk/wandb_init.py

@@ -273,8 +273,9 @@ class _WandbInit:
 
         tensorboard = kwargs.pop("tensorboard", None)
         sync_tensorboard = kwargs.pop("sync_tensorboard", None)
-        if tensorboard or sync_tensorboard and len(wandb.patched["tensorboard"]) == 0:
-            wandb.tensorboard.patch()  # type: ignore
+        if tensorboard or sync_tensorboard:
+            if len(wandb.patched["tensorboard"]) == 0:
+                wandb.tensorboard.patch()  # type: ignore
             with telemetry.context(obj=self._init_telemetry_obj) as tel:
                 tel.feature.tensorboard_sync = True
 
@@ -566,6 +567,7 @@ class _WandbInit:
             "watch",
             "unwatch",
             "upsert_artifact",
+            "_finish",
         ):
             setattr(drun, symbol, lambda *_, **__: None)  # type: ignore
         # attributes
@@ -736,7 +738,7 @@ class _WandbInit:
                 tel.feature.flow_control_disabled = True
             if self.settings._flow_control_custom:
                 tel.feature.flow_control_custom = True
-            if self.settings._require_core:
+            if not self.settings._require_legacy_service:
                 tel.feature.core = True
             if self.settings._shared:
                 wandb.termwarn(
@@ -1153,6 +1155,7 @@ def init(
             mode if a user isn't logged in to W&B. (default: `False`)
         sync_tensorboard: (bool, optional) Synchronize wandb logs from tensorboard or
             tensorboardX and save the relevant events file. (default: `False`)
+        tensorboard: (bool, optional) Alias for `sync_tensorboard`, deprecated.
         monitor_gym: (bool, optional) Automatically log videos of environment when
             using OpenAI Gym. (default: `False`)
             See [our guide to this integration](https://docs.wandb.com/guides/integrations/openai-gym).
@@ -1166,6 +1169,12 @@ def init(
             a moment in a previous run to fork a new run from. Creates a new run that picks up
             logging history from the specified run at the specified moment. The target run must
             be in the current project. Example: `fork_from="my-run-id?_step=1234"`.
+        resume_from: (str, optional) A string with the format {run_id}?_step={step} describing
+            a moment in a previous run to resume a run from. This allows users to truncate
+            the history logged to a run at an intermediate step and resume logging from that step.
+            It uses run forking under the hood. The target run must be in the
+            current project. Example: `resume_from="my-run-id?_step=1234"`.
+        settings: (dict, wandb.Settings, optional) Settings to use for this run. (default: None)
 
     Examples:
     ### Set where the run is logged

wandb/sdk/wandb_login.py

@@ -64,6 +64,7 @@ def login(
             "allow", only create an anonymous user if the user
             isn't already logged in. If set to "never", never log a
             user anonymously. Default set to "never".
+        key: (string, optional) The API key to use.
         relogin: (bool, optional) If true, will re-prompt for API key.
         host: (string, optional) The host to connect to.
         force: (bool, optional) If true, will force a relogin.

wandb/sdk/wandb_manager.py

@@ -13,7 +13,6 @@ import wandb
 from wandb import env, trigger
 from wandb.errors import Error
 from wandb.sdk.lib.exit_hooks import ExitHooks
-from wandb.sdk.lib.import_hooks import unregister_all_post_import_hooks
 
 if TYPE_CHECKING:
     from wandb.proto import wandb_settings_pb2
@@ -163,8 +162,6 @@ class _Manager:
         This sends a teardown record to the process. An exception is raised if
         the process has already been shut down.
         """
-        unregister_all_post_import_hooks()
-
         if self._atexit_lambda:
             atexit.unregister(self._atexit_lambda)
             self._atexit_lambda = None

wandb/sdk/wandb_require.py

@@ -13,7 +13,7 @@ import os
 from typing import Optional, Sequence, Union
 
 import wandb
-from wandb.env import _REQUIRE_CORE
+from wandb.env import _REQUIRE_LEGACY_SERVICE
 from wandb.errors import UnsupportedError
 from wandb.sdk import wandb_run
 from wandb.sdk.lib.wburls import wburls
@@ -41,7 +41,12 @@ class _Requires:
         self._require_service()
 
     def require_core(self) -> None:
-        os.environ[_REQUIRE_CORE] = "true"
+        wandb.termwarn(
+            "`wandb.require('core')` is redundant as it is now the default behavior."
+        )
+
+    def require_legacy_service(self) -> None:
+        os.environ[_REQUIRE_LEGACY_SERVICE] = "true"
 
     def apply(self) -> None:
         """Call require_* method for supported features."""

wandb/sdk/wandb_run.py

@@ -57,6 +57,7 @@ from wandb.sdk.lib.import_hooks import (
     unregister_post_import_hook,
 )
 from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath
+from wandb.sdk.lib.viz import CustomChart, Visualize, custom_chart
 from wandb.util import (
     _is_artifact_object,
     _is_artifact_string,
@@ -66,9 +67,9 @@ from wandb.util import (
     add_import_hook,
     parse_artifact_string,
 )
-from wandb.viz import CustomChart, Visualize, custom_chart
 
 from . import wandb_config, wandb_metric, wandb_summary
+from .artifacts._validators import validate_aliases, validate_tags
 from .data_types._dtypes import TypeRegistry
 from .interface.interface import FilesDict, GlobStr, InterfaceBase, PolicyName
 from .interface.summary_record import SummaryRecord
@@ -1654,7 +1655,8 @@ class Run:
             if os.getpid() != self._init_pid or self._is_attached:
                 wandb.termwarn(
                     "Note that setting step in multiprocessing can result in data loss. "
-                    "Please log your step values as a metric such as 'global_step'",
+                    "Please use `run.define_metric(...)` to define a custom metric "
+                    "to log your step values.",
                     repeat=False,
                 )
             # if step is passed in when tensorboard_sync is used we honor the step passed
@@ -1662,8 +1664,9 @@ class Run:
             # this history later on in publish_history()
             if len(wandb.patched["tensorboard"]) > 0:
                 wandb.termwarn(
-                    "Step cannot be set when using syncing with tensorboard. "
-                    "Please log your step values as a metric such as 'global_step'",
+                    "Step cannot be set when using tensorboard syncing. "
+                    "Please use `run.define_metric(...)` to define a custom metric "
+                    "to log your step values.",
                     repeat=False,
                 )
             if step > self._step:
@@ -2785,6 +2788,14 @@ class Run:
                 self,
             )
 
+        if (summary and "best" in summary) or goal is not None:
+            deprecate.deprecate(
+                deprecate.Deprecated.run__define_metric_best_goal,
+                "define_metric(summary='best', goal=...) is deprecated and will be removed. "
+                "Use define_metric(summary='min') or define_metric(summary='max') instead.",
+                self,
+            )
+
         return self._define_metric(
             name,
             step_metric,
@@ -3101,6 +3112,7 @@ class Run:
         name: Optional[str] = None,
         type: Optional[str] = None,
         aliases: Optional[List[str]] = None,
+        tags: Optional[List[str]] = None,
     ) -> Artifact:
         """Declare an artifact as an output of a run.
 
@@ -3121,12 +3133,17 @@ class Run:
             type: (str) The type of artifact to log, examples include `dataset`, `model`
             aliases: (list, optional) Aliases to apply to this artifact,
                 defaults to `["latest"]`
+            tags: (list, optional) Tags to apply to this artifact, if any.
 
         Returns:
             An `Artifact` object.
         """
         return self._log_artifact(
-            artifact_or_path, name=name, type=type, aliases=aliases
+            artifact_or_path,
+            name=name,
+            type=type,
+            aliases=aliases,
+            tags=tags,
         )
 
     @_run_decorator._noop_on_finish()
@@ -3243,6 +3260,7 @@ class Run:
         name: Optional[str] = None,
         type: Optional[str] = None,
         aliases: Optional[List[str]] = None,
+        tags: Optional[List[str]] = None,
         distributed_id: Optional[str] = None,
         finalize: bool = True,
         is_user_created: bool = False,
@@ -3253,13 +3271,17 @@ class Run:
             wandb.termwarn(
                 "Artifacts logged anonymously cannot be claimed and expire after 7 days."
             )
+
         if not finalize and distributed_id is None:
             raise TypeError("Must provide distributed_id if artifact is not finalize")
+
         if aliases is not None:
-            if any(invalid in alias for alias in aliases for invalid in ["/", ":"]):
-                raise ValueError(
-                    "Aliases must not contain any of the following characters: /, :"
-                )
+            aliases = validate_aliases(aliases)
+
+        # Check if artifact tags are supported
+        if tags is not None:
+            tags = validate_tags(tags)
+
         artifact, aliases = self._prepare_artifact(
             artifact_or_path, name, type, aliases
         )
@@ -3271,6 +3293,7 @@ class Run:
                     self,
                     artifact,
                     aliases,
+                    tags,
                     self.step,
                     finalize=finalize,
                     is_user_created=is_user_created,
@@ -3282,6 +3305,7 @@ class Run:
                     self,
                     artifact,
                     aliases,
+                    tags,
                     finalize=finalize,
                     is_user_created=is_user_created,
                     use_after_commit=use_after_commit,
@@ -3291,6 +3315,7 @@ class Run:
                 self,
                 artifact,
                 aliases,
+                tags,
                 finalize=finalize,
                 is_user_created=is_user_created,
                 use_after_commit=use_after_commit,
@@ -3364,6 +3389,7 @@ class Run:
                 "You must pass an instance of wandb.Artifact or a "
                 "valid file path to log_artifact"
             )
+
         artifact.finalize()
         return artifact, _resolve_aliases(aliases)
 
@@ -3502,7 +3528,7 @@ class Run:
             registered_model_name: (str) - the name of the registered model that the model is to be linked to.
                 A registered model is a collection of model versions linked to the model registry, typically representing a
                 team's specific ML Task. The entity that this registered model belongs to will be derived from the run
-                name: (str, optional) - the name of the model artifact that files in 'path' will be logged to. This will
+            name: (str, optional) - the name of the model artifact that files in 'path' will be logged to. This will
                 default to the basename of the path prepended with the current run id  if not specified.
             aliases: (List[str], optional) - alias(es) that will only be applied on this linked artifact
                 inside the registered model.
@@ -4234,13 +4260,13 @@ class Run:
         printer: Union["PrinterTerm", "PrinterJupyter"],
     ) -> None:
         """Prints a message advertising the upcoming core release."""
-        if quiet or settings._require_core:
+        if quiet or not settings._require_legacy_service:
             return
 
         printer.display(
-            "The new W&B backend becomes opt-out in version 0.18.0;"
-            ' try it out with `wandb.require("core")`!'
-            " See https://wandb.me/wandb-core for more information.",
+            "The legacy backend is deprecated. In future versions, `wandb-core` will become "
+            "the sole backend service, and the `wandb.require('legacy-service')` flag will be removed. "
+            "For more information, visit https://wandb.me/wandb-core",
             level="warn",
         )