wandb 0.18.5__py3-none-macosx_11_0_arm64.whl → 0.18.7__py3-none-macosx_11_0_arm64.whl

wandb/sdk/internal/file_stream.py

@@ -575,7 +575,7 @@ class FileStreamApi:
     def push(self, filename: str, data: str) -> None:
         """Push a chunk of a file to the streaming endpoint.
 
-        Arguments:
+        Args:
             filename: Name of file to append to.
             data: Text to append to the file.
         """
@@ -584,7 +584,7 @@ class FileStreamApi:
     def push_success(self, artifact_id: str, save_name: str) -> None:
         """Notification that a file upload has been successfully completed.
 
-        Arguments:
+        Args:
             artifact_id: ID of artifact
             save_name: saved name of the uploaded file
         """
@@ -595,7 +595,7 @@ class FileStreamApi:
 
         Anything pushed after finish will be dropped.
 
-        Arguments:
+        Args:
             exitcode: The exitcode of the watched process.
         """
         logger.info("file stream finish called")
@@ -620,7 +620,7 @@ def request_with_retry(
 ) -> Union["requests.Response", "requests.RequestException"]:
     """Perform a requests http call, retrying with exponential backoff.
 
-    Arguments:
+    Args:
         func:        An http-requesting function to call, like requests.post
         max_retries: Maximum retries before giving up.
                      By default, we retry 30 times in ~2 hours before dropping the chunk

wandb/sdk/internal/handler.py

@@ -21,6 +21,7 @@ from typing import (
     cast,
 )
 
+from wandb.errors.links import url_registry
 from wandb.proto.wandb_internal_pb2 import (
     HistoryRecord,
     InternalMessages,
@@ -37,7 +38,7 @@ from wandb.proto.wandb_internal_pb2 import (
 )
 
 from ..interface.interface_queue import InterfaceQueue
-from ..lib import handler_util, proto_util, tracelog, wburls
+from ..lib import handler_util, proto_util
 from . import context, sample, tb_watcher
 from .settings_static import SettingsStatic
 from .system.system_monitor import SystemMonitor
@@ -162,11 +163,9 @@ class HandleManager:
     def _dispatch_record(self, record: Record, always_send: bool = False) -> None:
         if always_send:
             record.control.always_send = True
-        tracelog.log_message_queue(record, self._writer_q)
         self._writer_q.put(record)
 
     def _respond_result(self, result: Result) -> None:
-        tracelog.log_message_queue(result, self._result_q)
         context_id = context.context_id_from_result(result)
         self._context_keeper.release(context_id)
         self._result_q.put(result)
@@ -578,7 +577,7 @@ class HandleManager:
                 if not self._dropped_history:
                     message = (
                         "Step only supports monotonically increasing values, use define_metric to set a custom x "
-                        f"axis. For details see: {wburls.wburls.get('wandb_define_metric')}"
+                        f"axis. For details see: {url_registry.url('define-metric')}"
                     )
                     self._internal_messages.warning.append(message)
                     self._dropped_history = True

wandb/sdk/internal/internal.py

@@ -28,7 +28,6 @@ import psutil
 import wandb
 
 from ..interface.interface_queue import InterfaceQueue
-from ..lib import tracelog
 from . import context, handler, internal_util, sender, writer
 
 if TYPE_CHECKING:
@@ -55,7 +54,7 @@ def wandb_internal(
 
     Read from record queue and dispatch work to various threads.
 
-    Arguments:
+    Args:
         settings: settings object
         record_q: records to be handled
         result_q: for sending results back
@@ -63,7 +62,6 @@ def wandb_internal(
     """
     # mark this process as internal
     wandb._set_internal_process()  # type: ignore
-    _setup_tracelog()
     started = time.time()
 
     # any sentry events in the internal process will be tagged as such
@@ -88,8 +86,6 @@ def wandb_internal(
         datetime.fromtimestamp(started),
     )
 
-    tracelog.annotate_queue(record_q, "record_q")
-    tracelog.annotate_queue(result_q, "result_q")
     publish_interface = InterfaceQueue(record_q=record_q)
 
     stopped = threading.Event()
@@ -98,10 +94,8 @@ def wandb_internal(
     context_keeper = context.ContextKeeper()
 
     send_record_q: Queue[Record] = queue.Queue()
-    tracelog.annotate_queue(send_record_q, "send_q")
 
     write_record_q: Queue[Record] = queue.Queue()
-    tracelog.annotate_queue(write_record_q, "write_q")
 
     record_sender_thread = SenderThread(
         settings=_settings,
@@ -184,14 +178,6 @@ def wandb_internal(
     close_internal_log()
 
 
-def _setup_tracelog() -> None:
-    # TODO: remove this temporary hack, need to find a better way to pass settings
-    # to the server.  for now lets just look at the environment variable we need
-    tracelog_mode = os.environ.get("WANDB_TRACELOG")
-    if tracelog_mode:
-        tracelog.enable(tracelog_mode)
-
-
 def configure_logging(
     log_fname: str, log_level: int, run_id: Optional[str] = None
 ) -> None:

wandb/sdk/internal/internal_api.py

@@ -22,6 +22,7 @@ from typing import (
     List,
     Mapping,
     MutableMapping,
+    NamedTuple,
     Optional,
     Sequence,
     TextIO,
@@ -169,6 +170,50 @@ class _ThreadLocalData(threading.local):
         self.context = None
 
 
+class _OrgNames(NamedTuple):
+    entity_name: str
+    display_name: str
+
+
+def _match_org_with_fetched_org_entities(
+    organization: str, orgs: Sequence[_OrgNames]
+) -> str:
+    """Match the organization provided in the path with the org entity or org name of the input entity.
+
+    Args:
+        organization: The organization name to match
+        orgs: List of tuples containing (org_entity_name, org_display_name)
+
+    Returns:
+        str: The matched org entity name
+
+    Raises:
+        ValueError: If no matching organization is found or if multiple orgs exist without a match
+    """
+    for org_names in orgs:
+        if organization in org_names:
+            wandb.termwarn(
+                "Registries can be linked/fetched using a shorthand form without specifying the organization name. "
+                "Try using shorthand path format: <my_registry_name>/<artifact_name> or "
+                "just <my_registry_name> if fetching just the project."
+            )
+            return org_names.entity_name
+
+    if len(orgs) == 1:
+        raise ValueError(
+            f"Expecting the organization name or entity name to match {orgs[0].display_name!r} "
+            f"and cannot be linked/fetched with {organization!r}. "
+            "Please update the target path with the correct organization name."
+        )
+
+    raise ValueError(
+        "Personal entity belongs to multiple organizations "
+        f"and cannot be linked/fetched with {organization!r}. "
+        "Please update the target path with the correct organization name "
+        "or use a team entity in the entity settings."
+    )
+
+
 class Api:
     """W&B Internal Api wrapper.
 
@@ -178,7 +223,7 @@ class Api:
         directory. If none can be found, we look in the current user's home
         directory.
 
-    Arguments:
+    Args:
         default_settings(dict, optional): If you aren't using a settings
         file, or you wish to override the section to use in the settings file
         Override the settings here.
@@ -313,6 +358,7 @@ class Api:
         self.server_create_artifact_input_info: Optional[List[str]] = None
         self.server_artifact_fields_info: Optional[List[str]] = None
         self.server_organization_type_fields_info: Optional[List[str]] = None
+        self.server_supports_enabling_artifact_usage_tracking: Optional[bool] = 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
@@ -434,7 +480,7 @@ class Api:
     def settings(self, key: Optional[str] = None, section: Optional[str] = None) -> Any:
         """The settings overridden from the wandb/settings file.
 
-        Arguments:
+        Args:
             key (str, optional): If provided only this setting is returned
             section (str, optional): If provided this section of the setting file is
             used, defaults to "default"
@@ -514,7 +560,7 @@ class Api:
     ) -> Tuple[str, str]:
         """Parse a slug into a project and run.
 
-        Arguments:
+        Args:
             slug (str): The slug to parse
             project (str, optional): The project to use, if not provided it will be
             inferred from the slug
@@ -953,7 +999,7 @@ class Api:
     def list_projects(self, entity: Optional[str] = None) -> List[Dict[str, str]]:
         """List projects in W&B scoped by entity.
 
-        Arguments:
+        Args:
             entity (str, optional): The entity to scope this project to.
 
         Returns:
@@ -985,7 +1031,7 @@ class Api:
     def project(self, project: str, entity: Optional[str] = None) -> "_Response":
         """Retrieve project.
 
-        Arguments:
+        Args:
             project (str): The project to get details for
             entity (str, optional): The entity to scope this project to.
 
@@ -1020,7 +1066,7 @@ class Api:
     ) -> Dict[str, Any]:
         """Retrieve sweep.
 
-        Arguments:
+        Args:
             sweep (str): The sweep to get details for
             specs (str): history specs
             project (str, optional): The project to scope this sweep to.
@@ -1093,7 +1139,7 @@ class Api:
     ) -> List[Dict[str, str]]:
         """List runs in W&B scoped by project.
 
-        Arguments:
+        Args:
             project (str): The project to scope the runs to
             entity (str, optional): The entity to scope this project to.  Defaults to public models
 
@@ -1134,7 +1180,7 @@ class Api:
     ) -> Tuple[str, Dict[str, Any], Optional[str], Dict[str, Any]]:
         """Get the relevant configs for a run.
 
-        Arguments:
+        Args:
             project (str): The project to download, (can include bucket)
             run (str, optional): The run to download
             entity (str, optional): The entity to scope this project to.
@@ -1223,7 +1269,7 @@ class Api:
     ) -> Optional[Dict[str, Any]]:
         """Check if a run exists and get resume information.
 
-        Arguments:
+        Args:
             entity (str): The entity to scope this project to.
             project_name (str): The project to download, (can include bucket)
             name (str): The run to download
@@ -1328,7 +1374,7 @@ class Api:
     ) -> Dict[str, Any]:
         """Create a new project.
 
-        Arguments:
+        Args:
             project (str): The project to create
             description (str, optional): A description of this project
             entity (str, optional): The entity to scope this project to.
@@ -2156,7 +2202,7 @@ class Api:
     ) -> Tuple[dict, bool, Optional[List]]:
         """Update a run.
 
-        Arguments:
+        Args:
             id (str, optional): The existing run to update
             name (str, optional): The name of the run to create
             group (str, optional): Name of the group this run is a part of
@@ -2343,7 +2389,7 @@ class Api:
     ) -> dict:
         """Rewinds a run to a previous state.
 
-        Arguments:
+        Args:
             run_name (str): The name of the run to rewind
             metric_name (str): The name of the metric to rewind to
             metric_value (float): The value of the metric to rewind to
@@ -2530,7 +2576,7 @@ class Api:
     ) -> Tuple[str, List[str], Dict[str, Dict[str, Any]]]:
         """Generate temporary resumable upload urls.
 
-        Arguments:
+        Args:
             project (str): The project to download
             files (list or dict): The filenames to upload
             run (str, optional): The run to upload to
@@ -2667,7 +2713,7 @@ class Api:
     ) -> Dict[str, Dict[str, str]]:
         """Generate download urls.
 
-        Arguments:
+        Args:
             project (str): The project to download
             run (str): The run to upload to
             entity (str, optional): The entity to scope this project to.  Defaults to wandb models
@@ -2726,7 +2772,7 @@ class Api:
     ) -> Optional[Dict[str, str]]:
         """Generate download urls.
 
-        Arguments:
+        Args:
             project (str): The project to download
             file_name (str): The name of the file to download
             run (str): The run to upload to
@@ -2779,7 +2825,7 @@ class Api:
     def download_file(self, url: str) -> Tuple[int, requests.Response]:
         """Initiate a streaming download.
 
-        Arguments:
+        Args:
             url (str): The url to download
 
         Returns:
@@ -2813,7 +2859,7 @@ class Api:
     ) -> Tuple[str, Optional[requests.Response]]:
         """Download a file from a run and write it to wandb/.
 
-        Arguments:
+        Args:
             metadata (obj): The metadata object for the file to download. Comes from Api.download_urls().
             out_dir (str, optional): The directory to write the file to. Defaults to wandb/
 
@@ -2877,7 +2923,7 @@ class Api:
     ) -> Optional[requests.Response]:
         """Upload a file chunk to S3 with failure resumption.
 
-        Arguments:
+        Args:
             url: The url to download
             upload_chunk: The path to the file you want to upload
             extra_headers: A dictionary of extra headers to send with the request
@@ -2930,7 +2976,7 @@ class Api:
     ) -> Optional[requests.Response]:
         """Upload a file to W&B with failure resumption.
 
-        Arguments:
+        Args:
             url: The url to download
             file: The path to the file you want to upload
             callback: A callback which is passed the number of
@@ -3002,7 +3048,7 @@ class Api:
     ) -> dict:
         """Register a new agent.
 
-        Arguments:
+        Args:
             host (str): hostname
             sweep_id (str): sweep id
             project_name: (str): model that contains sweep
@@ -3052,7 +3098,7 @@ class Api:
     ) -> List[Dict[str, Any]]:
         """Notify server about agent state, receive commands.
 
-        Arguments:
+        Args:
             agent_id (str): agent_id
             metrics (dict): system metrics
             run_states (dict): run_id: state mapping
@@ -3161,7 +3207,7 @@ class Api:
     ) -> Tuple[str, List[str]]:
         """Upsert a sweep object.
 
-        Arguments:
+        Args:
             config (dict): sweep config (will be converted to yaml)
             controller (str): controller to use
             launch_scheduler (str): launch scheduler to use
@@ -3342,7 +3388,7 @@ class Api:
     ) -> "List[requests.Response]":
         """Download files from W&B.
 
-        Arguments:
+        Args:
             project (str): The project to download
             run (str, optional): The run to upload to
             entity (str, optional): The entity to scope this project to.  Defaults to wandb models
@@ -3377,7 +3423,7 @@ class Api:
     ) -> "List[Optional[requests.Response]]":
         """Uploads multiple files to W&B.
 
-        Arguments:
+        Args:
             files (list or dict): The filenames to upload, when dict the values are open files
             run (str, optional): The run to upload to
             entity (str, optional): The entity to scope this project to.  Defaults to wandb models
@@ -3496,7 +3542,9 @@ class Api:
         org_entity = ""
         if is_artifact_registry_project(project):
             try:
-                org_entity = self._resolve_org_entity_name(entity, organization)
+                org_entity = self._resolve_org_entity_name(
+                    entity=entity, organization=organization
+                )
             except ValueError as e:
                 wandb.termerror(str(e))
                 raise
@@ -3538,47 +3586,67 @@ class Api:
         # the organization parameter, or an error if it is empty. Otherwise, this returns the
         # fetched value after validating that the given organization, if not empty, matches
         # either the org's display or entity name.
+
+        if not entity:
+            raise ValueError("Entity name is required to resolve org entity name.")
+
         org_fields = self.server_organization_type_introspection()
-        can_fetch_org_entity = "orgEntity" in org_fields
-        if not organization and not can_fetch_org_entity:
+        can_shorthand_org_entity = "orgEntity" in org_fields
+        if not organization and not can_shorthand_org_entity:
             raise ValueError(
                 "Fetching Registry artifacts without inputting an organization "
                 "is unavailable for your server version. "
                 "Please upgrade your server to 0.50.0 or later."
             )
-        if not can_fetch_org_entity:
+        if not can_shorthand_org_entity:
             # Server doesn't support fetching org entity to validate,
             # assume org entity is correctly inputted
             return organization
 
-        org_entity, org_name = self.fetch_org_entity_from_entity(entity)
+        orgs_from_entity = self._fetch_orgs_and_org_entities_from_entity(entity)
         if organization:
-            if organization != org_name and organization != org_entity:
-                raise ValueError(
-                    f"Artifact belongs to the organization {org_name!r} "
-                    f"and cannot be linked/fetched with {organization!r}. "
-                    "Please update the target path with the correct organization name."
-                )
-            wandb.termwarn(
-                "Registries can be linked/fetched using a shorthand form without specifying the organization name. "
-                "Try using shorthand path format: <my_registry_name>/<artifact_name>"
+            return _match_org_with_fetched_org_entities(organization, orgs_from_entity)
+
+        # If no input organization provided, error if entity belongs to multiple orgs because we
+        # cannot determine which one to use.
+        if len(orgs_from_entity) > 1:
+            raise ValueError(
+                f"Personal entity {entity!r} belongs to multiple organizations "
+                "and cannot be used without specifying the organization name. "
+                "Please specify the organization in the Registry path or use a team entity in the entity settings."
             )
-        return org_entity
+        return orgs_from_entity[0].entity_name
+
+    def _fetch_orgs_and_org_entities_from_entity(self, entity: str) -> List[_OrgNames]:
+        """Fetches organization entity names and display names for a given entity.
 
-    def fetch_org_entity_from_entity(self, entity: str) -> Tuple[str, str]:
+        Args:
+            entity (str): Entity name to lookup. Can be either a personal or team entity.
+
+        Returns:
+            List[_OrgNames]: List of _OrgNames tuples. (_OrgNames(entity_name, display_name))
+
+        Raises:
+        ValueError: If entity is not found, has no organizations, or other validation errors.
+        """
         query = gql(
             """
-            query FetchOrgEntityFromEntity(
-                $entityName: String!,
-            ) {
+            query FetchOrgEntityFromEntity($entityName: String!) {
                 entity(name: $entityName) {
-                    isTeam
                     organization {
                         name
                         orgEntity {
                             name
                         }
                     }
+                    user {
+                        organizations {
+                            name
+                            orgEntity {
+                                name
+                            }
+                        }
+                    }
                 }
             }
             """
@@ -3589,28 +3657,40 @@ class Api:
                 "entityName": entity,
             },
         )
-        try:
-            is_team = response["entity"].get("isTeam", False)
-            org = response["entity"]["organization"]
-            org_name = org["name"] or ""
-            org_entity_name = org["orgEntity"]["name"] or ""
-        except (LookupError, TypeError) as e:
-            if is_team:
-                # This path should pretty much never be reached as all team entities have an organization.
+
+        # Parse organization from response
+        entity_resp = response["entity"]["organization"]
+        user_resp = response["entity"]["user"]
+        # Check for organization under team/org entity type
+        if entity_resp:
+            org_name = entity_resp.get("name")
+            org_entity_name = entity_resp.get("orgEntity") and entity_resp[
+                "orgEntity"
+            ].get("name")
+            if not org_name or not org_entity_name:
                 raise ValueError(
-                    f"Unable to find an organization under entity {entity!r}. "
-                ) from e
-            else:
+                    f"Unable to find an organization under entity {entity!r}."
+                )
+            return [_OrgNames(entity_name=org_entity_name, display_name=org_name)]
+        # Check for organization under personal entity type, where a user can belong to multiple orgs
+        elif user_resp:
+            orgs = user_resp.get("organizations", [])
+            org_entities_return = [
+                _OrgNames(
+                    entity_name=org["orgEntity"]["name"], display_name=org["name"]
+                )
+                for org in orgs
+                if org.get("orgEntity") and org.get("name")
+            ]
+            if not org_entities_return:
                 raise ValueError(
-                    f"Unable to resolve an organization associated with the entity: {entity!r} "
-                    "that is initialized in the API or Run settings. This could be because "
-                    f"{entity!r} is a personal entity or the team entity doesn't exist. "
-                    "Please re-initialize the API or Run with a team entity using "
-                    "wandb.Api(overrides={'entity': '<my_team_entity>'}) "
-                    "or wandb.init(entity='<my_team_entity>') "
-                ) from e
+                    f"Unable to resolve an organization associated with personal entity: {entity!r}. "
+                    "This could be because its a personal entity that doesn't belong to any organizations. "
+                    "Please specify the organization in the Registry path or use a team entity in the entity settings."
+                )
+            return org_entities_return
         else:
-            return org_entity_name, org_name
+            raise ValueError(f"Unable to find an organization under entity {entity!r}.")
 
     def use_artifact(
         self,
@@ -3701,6 +3781,41 @@ class Api:
 
         return self.server_organization_type_fields_info
 
+    # Fetch input arguments for the "artifact" endpoint on the "Project" type
+    def server_project_type_introspection(self) -> bool:
+        if self.server_supports_enabling_artifact_usage_tracking is not None:
+            return self.server_supports_enabling_artifact_usage_tracking
+
+        query_string = """
+            query ProbeServerProjectInfo {
+                ProjectInfoType: __type(name:"Project") {
+                    fields {
+                        name
+                        args {
+                            name
+                        }
+                    }
+                }
+            }
+        """
+
+        query = gql(query_string)
+        res = self.gql(query)
+        input_fields = res.get("ProjectInfoType", {}).get("fields", [{}])
+        artifact_args: List[Dict[str, str]] = next(
+            (
+                field.get("args", [])
+                for field in input_fields
+                if field.get("name") == "artifact"
+            ),
+            [],
+        )
+        self.server_supports_enabling_artifact_usage_tracking = any(
+            arg.get("name") == "enableTracking" for arg in artifact_args
+        )
+
+        return self.server_supports_enabling_artifact_usage_tracking
+
     def create_artifact_type(
         self,
         artifact_type_name: str,

wandb/sdk/internal/internal_util.py

@@ -11,8 +11,6 @@ import threading
 import time
 from typing import TYPE_CHECKING, Optional, Tuple, Type, Union
 
-from ..lib import tracelog
-
 if TYPE_CHECKING:
     from queue import Queue
     from threading import Event
@@ -95,6 +93,5 @@ class RecordLoopThread(ExceptionThread):
                 record = self._input_record_q.get(timeout=1)
             except queue.Empty:
                 continue
-            tracelog.log_message_dequeue(record, self._input_record_q)
             self._process(record)
         self._finish()

wandb/sdk/internal/job_builder.py

@@ -313,7 +313,8 @@ class JobBuilder:
             "build_context": metadata.get("build_context"),
             "dockerfile": metadata.get("dockerfile"),
         }
-        name = self._make_job_name(self._logged_code_artifact["name"])
+        artifact_basename, *_ = self._logged_code_artifact["name"].split(":")
+        name = self._make_job_name(artifact_basename)
 
         return source, name
 
@@ -380,7 +381,7 @@ class JobBuilder:
     ]:
         """Construct a job source dict and name from the current run.
 
-        Arguments:
+        Args:
             source_type (str): The type of source to build the job from. One of
                 "repo", "artifact", or "image".
         """
@@ -427,7 +428,7 @@ class JobBuilder:
     ) -> Optional[Artifact]:
         """Build a job artifact from the current run.
 
-        Arguments:
+        Args:
             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

wandb/sdk/internal/sender.py

@@ -53,7 +53,6 @@ from wandb.sdk.lib import (
     proto_util,
     redirect,
     telemetry,
-    tracelog,
 )
 from wandb.sdk.lib.mailbox import ContextCancelledError
 from wandb.sdk.lib.proto_util import message_to_dict
@@ -418,7 +417,6 @@ class SendManager:
         send_handler(record)
 
     def _respond_result(self, result: "Result") -> None:
-        tracelog.log_message_queue(result, self._result_q)
         context_id = context.context_id_from_result(result)
         self._context_keeper.release(context_id)
         self._result_q.put(result)