wandb 0.18.4__py3-none-any.whl → 0.18.6__py3-none-any.whl

wandb/sdk/internal/internal_api.py

@@ -20,9 +20,9 @@ from typing import (
     Dict,
     Iterable,
     List,
-    Literal,
     Mapping,
     MutableMapping,
+    NamedTuple,
     Optional,
     Sequence,
     TextIO,
@@ -30,6 +30,11 @@ from typing import (
     Union,
 )
 
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
 import click
 import requests
 import yaml
@@ -165,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.
 
@@ -174,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.
@@ -309,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
@@ -430,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"
@@ -510,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
@@ -949,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:
@@ -981,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.
 
@@ -1016,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.
@@ -1089,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
 
@@ -1130,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.
@@ -1219,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
@@ -1324,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.
@@ -2152,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
@@ -2339,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
@@ -2526,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
@@ -2663,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
@@ -2722,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
@@ -2775,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:
@@ -2809,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/
 
@@ -2873,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
@@ -2926,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
@@ -2998,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
@@ -3048,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
@@ -3157,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
@@ -3338,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
@@ -3373,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
@@ -3492,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
@@ -3534,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_org_entity_from_entity(self, entity: str) -> Tuple[str, str]:
+    def _fetch_orgs_and_org_entities_from_entity(self, entity: str) -> List[_OrgNames]:
+        """Fetches organization entity names and display names for a given entity.
+
+        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
+                            }
+                        }
+                    }
                 }
             }
             """
@@ -3585,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,
@@ -3697,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/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/system/assets/__init__.py

@@ -4,7 +4,6 @@ __all__ = (
     "Disk",
     "GPU",
     "GPUAMD",
-    "GPUApple",
     "IPU",
     "Memory",
     "Network",
@@ -18,7 +17,6 @@ from .cpu import CPU
 from .disk import Disk
 from .gpu import GPU
 from .gpu_amd import GPUAMD
-from .gpu_apple import GPUApple
 from .ipu import IPU
 from .memory import Memory
 from .network import Network

wandb/sdk/internal/tb_watcher.py

@@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import wandb
 from wandb import util
-from wandb.plot.viz import CustomChart
+from wandb.plot import CustomChart
 from wandb.sdk.interface.interface import GlobStr
 from wandb.sdk.lib import filesystem
 
@@ -73,7 +73,7 @@ def is_tfevents_file_created_by(
     if not path:
         raise ValueError("Path must be a nonempty string")
     basename = os.path.basename(path)
-    if basename.endswith(".profile-empty") or basename.endswith(".sagemaker-uploaded"):
+    if basename.endswith((".profile-empty", ".sagemaker-uploaded")):
         return False
     fname_components = basename.split(".")
     try:
@@ -439,18 +439,19 @@ class TBEventConsumer:
 
     def _save_row(self, row: "HistoryDict") -> None:
         chart_keys = set()
-        for k in row:
-            if isinstance(row[k], CustomChart):
+        for k, v in row.items():
+            if isinstance(v, CustomChart):
                 chart_keys.add(k)
-                key = row[k].get_config_key(k)
-                value = row[k].get_config_value(
-                    "Vega2", row[k].user_query(f"{k}_table")
+                v.set_key(k)
+                self._tbwatcher._interface.publish_config(
+                    key=v.spec.config_key,
+                    val=v.spec.config_value,
                 )
-                row[k] = row[k]._data
-                self._tbwatcher._interface.publish_config(val=value, key=key)
 
         for k in chart_keys:
-            row[f"{k}_table"] = row.pop(k)
+            chart = row.pop(k)
+            if isinstance(chart, CustomChart):
+                row[chart.spec.table_key] = chart.table
 
         self._tbwatcher._interface.publish_history(
             row, run=self._internal_run, publish_step=False

wandb/sdk/launch/_launch.py

@@ -120,10 +120,11 @@ def resolve_agent_config(  # noqa: C901
         if isinstance(resolved_config.get("queue"), str):
             resolved_config["queues"].append(resolved_config["queue"])
         else:
-            raise LaunchError(
-                f"Invalid launch agent config for key 'queue' with type: {type(resolved_config.get('queue'))}"
-                + " (expected str). Specify multiple queues with the 'queues' key"
+            msg = (
+                "Invalid launch agent config for key 'queue' with type: {type(resolved_config.get('queue'))} "
+                "(expected str). Specify multiple queues with the 'queues' key"
             )
+            raise LaunchError(msg)
 
     keys = ["entity"]
     settings = {

wandb/sdk/launch/_launch_add.py

@@ -61,7 +61,7 @@ def launch_add(
         config: A dictionary containing the configuration for the run. May also contain
             resource specific arguments under the key "resource_args"
         template_variables: A dictionary containing values of template variables for a run queue.
-            Expected format of {"VAR_NAME": VAR_VALUE}
+            Expected format of `{"VAR_NAME": VAR_VALUE}`
         project: Target project to send launched run to
         entity: Target entity to send launched run to
         queue: the name of the queue to enqueue the run to
@@ -240,7 +240,7 @@ def _launch_add(
     public_api = public.Api()
     if job is not None:
         try:
-            public_api.artifact(job, type="job")
+            public_api._artifact(job, type="job")
         except (ValueError, CommError) as e:
             raise LaunchError(f"Unable to fetch job with name {job}: {e}")
 

wandb/sdk/launch/builder/kaniko_builder.py

@@ -199,7 +199,6 @@ class KanikoBuilder(AbstractBuilder):
 
     def login(self) -> None:
         """Login to the registry."""
-        pass
 
     async def _create_docker_ecr_config_map(
         self, job_name: str, corev1_client: client.CoreV1Api, repository: str

wandb/sdk/launch/create_job.py

@@ -168,6 +168,7 @@ def _create_job(
         return None, "", []
 
     job_builder = _configure_job_builder_for_partial(tempdir.name, job_source=job_type)
+    job_builder._settings.update(job_name=name)
     if job_type == "code":
         assert entrypoint is not None
         job_name = _make_code_artifact(

wandb/sdk/launch/environment/local_environment.py

@@ -12,7 +12,6 @@ class LocalEnvironment(AbstractEnvironment):
 
     def __init__(self) -> None:
         """Initialize a local environment by doing nothing."""
-        pass
 
     @classmethod
     def from_config(

wandb/sdk/launch/errors.py

@@ -4,16 +4,10 @@ from wandb.errors import Error
 class LaunchError(Error):
     """Raised when a known error occurs in wandb launch."""
 
-    pass
-
 
 class LaunchDockerError(Error):
     """Raised when Docker daemon is not running."""
 
-    pass
-
 
 class ExecutionError(Error):
     """Generic execution exception."""
-
-    pass

wandb/sdk/launch/registry/local_registry.py

@@ -20,7 +20,6 @@ class LocalRegistry(AbstractRegistry):
 
     def __init__(self) -> None:
         """Initialize a local registry."""
-        pass
 
     @classmethod
     def from_config(
@@ -40,7 +39,6 @@ class LocalRegistry(AbstractRegistry):
 
     async def verify(self) -> None:
         """Verify the local registry by doing nothing."""
-        pass
 
     async def get_username_password(self) -> Tuple[str, str]:
         """Get the username and password of the local registry."""