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

wandb/sdk/internal/internal_api.py

@@ -20,6 +20,7 @@ from typing import (
     Dict,
     Iterable,
     List,
+    Literal,
     Mapping,
     MutableMapping,
     Optional,
@@ -41,6 +42,7 @@ from wandb.apis.normalize import normalize_exceptions, parse_backend_error_messa
 from wandb.errors import AuthenticationError, CommError, UnsupportedError, UsageError
 from wandb.integration.sagemaker import parse_sm_secrets
 from wandb.old.settings import Settings
+from wandb.sdk.artifacts._validators import is_artifact_registry_project
 from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
 from wandb.sdk.lib.gql_request import GraphQLSession
 from wandb.sdk.lib.hashutil import B64MD5, md5_file_b64
@@ -306,6 +308,7 @@ class Api:
         self.server_use_artifact_input_info: Optional[List[str]] = None
         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._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
@@ -3366,7 +3369,7 @@ class Api:
         project: Optional[str] = None,
         description: Optional[str] = None,
         force: bool = True,
-        progress: Union[TextIO, bool] = False,
+        progress: Union[TextIO, Literal[False]] = False,
     ) -> "List[Optional[requests.Response]]":
         """Uploads multiple files to W&B.
 
@@ -3378,7 +3381,7 @@ class Api:
             description (str, optional): The description of the changes
             force (bool, optional): Whether to prevent push if git has uncommitted changes
             progress (callable, or stream): If callable, will be called with (chunk_bytes,
-                total_bytes) as argument else if True, renders a progress bar to stream.
+                total_bytes) as argument. If TextIO, renders a progress bar to it.
 
         Returns:
             A list of `requests.Response` objects
@@ -3439,8 +3442,8 @@ class Api:
                     )
                 else:
                     length = os.fstat(open_file.fileno()).st_size
-                    with click.progressbar(
-                        file=progress,  # type: ignore
+                    with click.progressbar(  # type: ignore
+                        file=progress,
                         length=length,
                         label=f"Uploading file: {file_name}",
                         fill_char=click.style("&", fg="green"),
@@ -3464,6 +3467,7 @@ class Api:
         entity: str,
         project: str,
         aliases: Sequence[str],
+        organization: str,
     ) -> Dict[str, Any]:
         template = """
                 mutation LinkArtifact(
@@ -3485,6 +3489,14 @@ class Api:
                     }
             """
 
+        org_entity = ""
+        if is_artifact_registry_project(project):
+            try:
+                org_entity = self._resolve_org_entity_name(entity, organization)
+            except ValueError as e:
+                wandb.termerror(str(e))
+                raise
+
         def replace(a: str, b: str) -> None:
             nonlocal template
             template = template.replace(a, b)
@@ -3500,7 +3512,7 @@ class Api:
             "clientID": client_id,
             "artifactID": server_id,
             "artifactPortfolioName": portfolio_name,
-            "entityName": entity,
+            "entityName": org_entity or entity,
             "projectName": project,
             "aliases": [
                 {"alias": alias, "artifactCollectionName": portfolio_name}
@@ -3513,6 +3525,89 @@ class Api:
         link_artifact: Dict[str, Any] = response["linkArtifact"]
         return link_artifact
 
+    def _resolve_org_entity_name(self, entity: str, organization: str = "") -> str:
+        # resolveOrgEntityName fetches the portfolio's org entity's name.
+        #
+        # The organization parameter may be empty, an org's display name, or an org entity name.
+        #
+        # If the server doesn't support fetching the org name of a portfolio, then this returns
+        # 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.
+        org_fields = self.server_organization_type_introspection()
+        can_fetch_org_entity = "orgEntity" in org_fields
+        if not organization and not can_fetch_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:
+            # 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)
+        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 org_entity
+
+    def fetch_org_entity_from_entity(self, entity: str) -> Tuple[str, str]:
+        query = gql(
+            """
+            query FetchOrgEntityFromEntity(
+                $entityName: String!,
+            ) {
+                entity(name: $entityName) {
+                    isTeam
+                    organization {
+                        name
+                        orgEntity {
+                            name
+                        }
+                    }
+                }
+            }
+            """
+        )
+        response = self.gql(
+            query,
+            variable_values={
+                "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.
+                raise ValueError(
+                    f"Unable to find an organization under entity {entity!r}. "
+                ) from e
+            else:
+                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
+        else:
+            return org_entity_name, org_name
+
     def use_artifact(
         self,
         artifact_id: str,
@@ -3580,6 +3675,28 @@ class Api:
             return artifact
         return None
 
+    # Fetch fields available in backend of Organization type
+    def server_organization_type_introspection(self) -> List[str]:
+        query_string = """
+            query ProbeServerOrganization {
+                OrganizationInfoType: __type(name:"Organization") {
+                    fields {
+                        name
+                    }
+                }
+            }
+        """
+
+        if self.server_organization_type_fields_info is None:
+            query = gql(query_string)
+            res = self.gql(query)
+            input_fields = res.get("OrganizationInfoType", {}).get("fields", [{}])
+            self.server_organization_type_fields_info = [
+                field["name"] for field in input_fields if "name" in field
+            ]
+
+        return self.server_organization_type_fields_info
+
     def create_artifact_type(
         self,
         artifact_type_name: str,

wandb/sdk/internal/sender.py

@@ -1455,16 +1455,29 @@ class SendManager:
         entity = link.portfolio_entity
         project = link.portfolio_project
         aliases = link.portfolio_aliases
+        organization = link.portfolio_organization
         logger.debug(
-            f"link_artifact params - client_id={client_id}, server_id={server_id}, pfolio={portfolio_name}, entity={entity}, project={project}"
+            f"link_artifact params - client_id={client_id}, server_id={server_id}, "
+            f"portfolio_name={portfolio_name}, entity={entity}, project={project}, "
+            f"organization={organization}"
         )
         if (client_id or server_id) and portfolio_name and entity and project:
             try:
                 self._api.link_artifact(
-                    client_id, server_id, portfolio_name, entity, project, aliases
+                    client_id,
+                    server_id,
+                    portfolio_name,
+                    entity,
+                    project,
+                    aliases,
+                    organization,
                 )
             except Exception as e:
-                result.response.log_artifact_response.error_message = f'error linking artifact to "{entity}/{project}/{portfolio_name}"; error: {e}'
+                org_or_entity = organization or entity
+                result.response.log_artifact_response.error_message = (
+                    f"error linking artifact to "
+                    f'"{org_or_entity}/{project}/{portfolio_name}"; error: {e}'
+                )
                 logger.warning("Failed to link artifact to portfolio: %s", e)
         self._respond_result(result)
 

wandb/sdk/launch/inputs/internal.py

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

wandb/sdk/lib/module.py

@@ -19,6 +19,8 @@ def set_global(
     log_model=None,
     use_model=None,
     link_model=None,
+    watch=None,
+    unwatch=None,
 ):
     if run:
         wandb.run = run
@@ -48,6 +50,10 @@ def set_global(
         wandb.use_model = use_model
     if link_model:
         wandb.link_model = link_model
+    if watch:
+        wandb.watch = watch
+    if unwatch:
+        wandb.unwatch = unwatch
 
 
 def unset_globals():
@@ -55,6 +61,12 @@ def unset_globals():
     wandb.config = preinit.PreInitObject("wandb.config")
     wandb.summary = preinit.PreInitObject("wandb.summary")
     wandb.log = preinit.PreInitCallable("wandb.log", wandb.wandb_sdk.wandb_run.Run.log)
+    wandb.watch = preinit.PreInitCallable(
+        "wandb.watch", wandb.wandb_sdk.wandb_run.Run.watch
+    )
+    wandb.unwatch = preinit.PreInitCallable(
+        "wandb.unwatch", wandb.wandb_sdk.wandb_run.Run.unwatch
+    )
     wandb.save = preinit.PreInitCallable(
         "wandb.save", wandb.wandb_sdk.wandb_run.Run.save
     )