o -wi5@sddlmZmZmZmZmZmZddlmZddl Z ddl m Z ddl m Z mZddlmZddlmZmZmZmZmZmZmZdd lmZdd lmZmZmZm Z dd l!m"Z"m#Z#erbdd lm$Z$Gd ddZ%dS)) TYPE_CHECKINGAnyDictListLiteralOptional)gqlN) ServerFeature)REGISTRY_PREFIXvalidate_project_name)Api)DELETE_PROJECT_GQLFETCH_REGISTRY_GQLRENAME_PROJECT_GQLUPSERT_REGISTRY_PROJECT_GQL DeleteProject RenameProjectUpsertRegistryProject)AddOnlyArtifactTypesList)"fetch_org_entity_from_organizationformat_gql_artifact_types_inputgql_to_registry_visibilityregistry_visibility_to_gql) CollectionsVersions)Clientc@seZdZdZ dd4d5Z!d>d6d7Z"d>d8d9Z#d efd:d;Z$dS)?Registryz"A single registry in the Registry.Nclientr organizationentitynameattrscCs8||_||_||_||_||_|dur||dSdSN)r_name _saved_name_entity _organization_update_attributes)selfrrr r!r"r*b/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/wandb/apis/public/registries/registry.py__init__#szRegistry.__init__returncCs|dd|_|jdurtd|jd|dd|_|dd|_td d |d id gD|_|d d|_|dd|_ t |dd|_ dS)z>Helper method to update instance attributes from a dictionary.idNz Registry z's id is not found descriptionallowAllArtifactTypesInRegistryFcss|] }|ddVqdS)noder!Nr*).0tr*r*r+ =s  z.Registry._update_attributes.. artifactTypesedges createdAt updatedAtaccess) get_id ValueErrorr! _description_allow_all_artifact_typesr_artifact_types _created_at _updated_atr _visibility)r)r"r*r*r+r(3s  zRegistry._update_attributescCs d|jS)zAFull name of the registry including the `wandb-registry-` prefix.zwandb-registry-)r!r)r*r*r+ full_nameDs zRegistry.full_namecC|jS)z:Name of the registry without the `wandb-registry-` prefix.r$rDr*r*r+r!Iz Registry.namevaluecCs ||_dSr#rGr)rIr*r*r+r!Ns cCrF)z$Organization entity of the registry.)r&rDr*r*r+r RrHzRegistry.entitycCrF)z"Organization name of the registry.)r'rDr*r*r+rWrHzRegistry.organizationcCrF)zDescription of the registry.r>rDr*r*r+r0\rHzRegistry.descriptioncC ||_dS)z$Set the description of the registry.NrKrJr*r*r+r0a cCrF)zReturns whether all artifact types are allowed in the registry. If `True` then artifacts of any type can be added to this registry. If `False` then artifacts are restricted to the types in `artifact_types` for this registry. r?rDr*r*r+allow_all_artifact_typesfsz!Registry.allow_all_artifact_typescCrL)z;Set whether all artifact types are allowed in the registry.NrNrJr*r*r+rOorMcCrF)a6Returns the artifact types allowed in the registry. If `allow_all_artifact_types` is `True` then `artifact_types` reflects the types previously saved or currently used in the registry. If `allow_all_artifact_types` is `False` then artifacts are restricted to the types in `artifact_types`. Note: Previously saved artifact types cannot be removed. Example: ```python registry.artifact_types.append("model") registry.save() # once saved, the artifact type `model` cannot be removed registry.artifact_types.append("accidentally_added") registry.artifact_types.remove( "accidentally_added" ) # Types can only be removed if it has not been saved yet ``` )r@rDr*r*r+artifact_typestszRegistry.artifact_typescCrF)z+Timestamp of when the registry was created.)rArDr*r*r+ created_atrHzRegistry.created_atcCrF)z0Timestamp of when the registry was last updated.)rBrDr*r*r+ updated_atrHzRegistry.updated_atcCs |j|jgSr#)r rErDr*r*r+paths z Registry.path)r restrictedcCrF)aVisibility of the registry. Returns: Literal["organization", "restricted"]: The visibility level. - "organization": Anyone in the organization can view this registry. You can edit their roles later from the settings in the UI. - "restricted": Only invited members via the UI can access this registry. Public sharing is disabled. rCrDr*r*r+ visibilitys zRegistry.visibilitycCrL)aSet the visibility of the registry. Args: value: The visibility level. Options are: - "organization": Anyone in the organization can view this registry. You can edit their roles later from the settings in the UI. - "restricted": Only invited members via the UI can access this registry. Public sharing is disabled. NrUrJr*r*r+rVs filtercCsd|ji}t|j|j||S)z2Returns the collections belonging to the registry.r!)rErrrr)rWregistry_filterr*r*r+ collectionsszRegistry.collectionscCsd|ji}t|j|j|d|S)z/Returns the versions belonging to the registry.r!N)rErrrrXr*r*r+versionsszRegistry.versionsrVr0rPc Cst||}t|}t|g} |rt|} t|} d|d|d} z|jtt|||| | | dd} Wn ty@t | w| ddsKt | t ||||| ddS) a%Create a new registry. The registry name must be unique within the organization. This function should be called using `api.create_registry()` Args: client: The GraphQL client. organization: The name of the organization. name: The name of the registry (without the `wandb-registry-` prefix). visibility: The visibility level ('organization' or 'restricted'). description: An optional description for the registry. artifact_types: An optional list of allowed artifact types. Returns: Registry: The newly created Registry object. Raises: ValueError: If a registry with the same name already exists in the organization or if the creation fails. zFailed to create registry  in organization .r0 entityNamer!r:r1r6variable_values upsertModelinsertedproject) rr r rrexecuterr Exceptionr=r) clsrrr!rVr0rP org_entityrEaccepted_artifact_typesvisibility_valueregistry_creation_errorresponser*r*r+creates@    zRegistry.createcCstz|jjttd|jid}t|}Wnty'td|j d|j w|j j s8td|j d|j dS)z*Delete the registry. This is irreversible.r.r`zFailed to delete registry:  in organization: N) rrerr r<rmodel_validaterfr=r!r delete_modelsuccess)r)rlresultr*r*r+deletes zRegistry.deletecCsd|jd|jd}z|jjtt|j|jdd}Wn ty't |w|ddur2t ||dd|_ |j durBt || |j dS) zPLoad the registry attributes from the backend to reflect the latest saved state.zFailed to load registry r\r])r!r_r`r Nrd) r!rrrerrrEr rfr=r"r()r)load_failure_messagerlr*r*r+loads(    z Registry.loadc Cs`ttjs td|rtdddddt|jt |j }t |j j }d|jd|j}t|j}z|jjtt|j|j|||j|d d }t|}Wn tyat|w|jjrttd |jd |jd | |dd|j|jkr|jjtt!|j||jdd }t"|}|j|_|j#jrtd|jd dSdSdS)z(Save registry attributes to the backend.zsaving the registry is not enabled on this wandb server version. Please upgrade your server version or contact support at support@wandb.com.zACannot update artifact types when `allows_all_artifact_types` is Tz . Set it to Fz first.z$Failed to save and update registry: rnr^r`zCreated registry r\z on saverbrd)r_oldProjectNamenewProjectNamezCreated new registry N)$ InternalApi_server_supportsr +INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION RuntimeError_no_updating_registry_typesr=r rErrVrrPdraftr!rr r%rrerrr0r rOrrorf upsert_modelrcwandbtermlogr(rrrename_project)r)rjnewly_added_typesregistry_save_errorfull_saved_namerlrrr*r*r+save*sb       z Registry.savecCst|jjdko |jS)Nr)lenrPr}rOrDr*r*r+r|dsz$Registry._no_updating_registry_typesr#)NN)r-N)%__name__ __module__ __qualname____doc__strrrrr,r(propertyrEr!setterr rr0rOboolrrPrQrRrSrrVrrZrr[ classmethodrrmrsrurr|r*r*r*r+r s       @  :r)&typingrrrrrr wandb_gqlrrwandb.proto.wandb_internal_pb2r wandb.sdk.artifacts._validatorsr r wandb.sdk.internal.internal_apir rxwandb.sdk.projects._generatedr rrrrrr_freezable_listr_utilsrrrrregistries_searchrrrrr*r*r*r+s   $