Skip to content

Commit 1aab402

Browse files
md-mqmd-mq
committed
Add async organization client
 * Add AsyncOrganizationClient with async decorator handling and awaited organization, team, run, version, and lineage SDK calls. * Refactor tests to delete base path directly in the patch settings utils function
1 parent c95a5dc commit 1aab402

5 files changed

Lines changed: 884 additions & 19 deletions

File tree

‎cli/polyaxon/_client/organization.py‎

Lines changed: 375 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,11 @@
33
from clipped.utils.query_params import get_query_params
44

55
from polyaxon._client.client import PolyaxonClient
6-
from polyaxon._client.decorators import client_handler, get_global_or_inline_config
6+
from polyaxon._client.decorators import (
7+
async_client_handler,
8+
client_handler,
9+
get_global_or_inline_config,
10+
)
711
from polyaxon._client.mixin import ClientMixin
812
from polyaxon._constants.globals import DEFAULT
913
from polyaxon._env_vars.getters.user import get_local_owner
@@ -115,7 +119,7 @@ def __init__(
115119
raise PolyaxonClientException("Please provide a valid owner.")
116120

117121
owner, team = split_owner_team_space(owner)
118-
self._client = client
122+
self._set_client(client)
119123
self._owner = owner or DEFAULT
120124
self._team = team
121125
self._organization_data = V1Organization.model_construct()
@@ -751,3 +755,372 @@ def list_runs_artifacts_lineage(
751755
return self.client.organizations_v1.get_organization_runs_artifacts_lineage(
752756
self.owner, **params
753757
)
758+
759+
760+
class AsyncOrganizationClient(OrganizationClient):
761+
_IS_ASYNC = True
762+
763+
@async_client_handler(check_no_op=True, check_offline=True)
764+
async def refresh_data(self):
765+
self._organization_data = await self.client.organizations_v1.get_organization(
766+
self.owner
767+
)
768+
if self._organization_data.name is None:
769+
self._organization_data.name = self.owner
770+
771+
@async_client_handler(check_no_op=True, check_offline=True)
772+
async def list(
773+
self,
774+
query: Optional[str] = None,
775+
sort: Optional[str] = None,
776+
limit: Optional[int] = None,
777+
offset: Optional[int] = None,
778+
) -> V1ListOrganizationsResponse:
779+
params = get_query_params(limit=limit, offset=offset, query=query, sort=sort)
780+
return await self.client.organizations_v1.list_organizations(**params)
781+
782+
@async_client_handler(check_no_op=True, check_offline=True)
783+
async def list_members(
784+
self,
785+
query: Optional[str] = None,
786+
sort: Optional[str] = None,
787+
limit: Optional[int] = None,
788+
offset: Optional[int] = None,
789+
) -> V1ListOrganizationMembersResponse:
790+
params = get_query_params(
791+
limit=limit or 20, offset=offset, query=query, sort=sort
792+
)
793+
return await self.client.organizations_v1.list_organization_members(
794+
self.owner, **params
795+
)
796+
797+
@async_client_handler(check_no_op=True, check_offline=True)
798+
async def get_member(self, user: str) -> V1OrganizationMember:
799+
return await self.client.organizations_v1.get_organization_member(
800+
self.owner, user
801+
)
802+
803+
@async_client_handler(check_no_op=True, check_offline=True)
804+
async def create_member(
805+
self,
806+
data: Union[Dict, V1OrganizationMember],
807+
email: Optional[str] = None,
808+
) -> V1OrganizationMember:
809+
return await self.client.organizations_v1.create_organization_member(
810+
self.owner,
811+
body=data,
812+
email=email,
813+
)
814+
815+
@async_client_handler(check_no_op=True, check_offline=True)
816+
async def update_member(
817+
self,
818+
user: str,
819+
data: Union[Dict, V1OrganizationMember],
820+
) -> V1OrganizationMember:
821+
return await self.client.organizations_v1.update_organization_member(
822+
self.owner,
823+
user,
824+
body=data,
825+
)
826+
827+
@async_client_handler(check_no_op=True, check_offline=True)
828+
async def patch_member(
829+
self,
830+
user: str,
831+
data: Union[Dict, V1OrganizationMember],
832+
) -> V1OrganizationMember:
833+
return await self.client.organizations_v1.patch_organization_member(
834+
self.owner,
835+
user,
836+
body=data,
837+
)
838+
839+
@async_client_handler(check_no_op=True, check_offline=True)
840+
async def delete_member(self, user: str):
841+
return await self.client.organizations_v1.delete_organization_member(
842+
self.owner, user
843+
)
844+
845+
@async_client_handler(check_no_op=True, check_offline=True)
846+
async def list_teams(
847+
self,
848+
query: Optional[str] = None,
849+
sort: Optional[str] = None,
850+
limit: Optional[int] = None,
851+
offset: Optional[int] = None,
852+
bookmarks: bool = False,
853+
mode: Optional[str] = None,
854+
) -> V1ListTeamsResponse:
855+
params = get_query_params(
856+
limit=limit or 20, offset=offset, query=query, sort=sort
857+
)
858+
if bookmarks:
859+
params["bookmarks"] = bookmarks
860+
if mode:
861+
params["mode"] = mode
862+
return await self.client.teams_v1.list_teams(self.owner, **params)
863+
864+
@async_client_handler(check_no_op=True, check_offline=True)
865+
async def list_runs(
866+
self,
867+
query: Optional[str] = None,
868+
sort: Optional[str] = None,
869+
limit: Optional[int] = None,
870+
offset: Optional[int] = None,
871+
) -> V1ListRunsResponse:
872+
params = get_query_params(
873+
limit=limit or 20, offset=offset, query=query, sort=sort
874+
)
875+
if self.team:
876+
return await self.client.teams_v1.get_team_runs(
877+
self.owner, self.team, **params
878+
)
879+
return await self.client.organizations_v1.get_organization_runs(
880+
self.owner, **params
881+
)
882+
883+
@async_client_handler(check_no_op=True, check_offline=True)
884+
async def get_run(self, uuid: str) -> V1Run:
885+
if self.team:
886+
return await self.client.teams_v1.get_team_run(
887+
self.owner, self.team, uuid
888+
)
889+
return await self.client.organizations_v1.get_organization_run(
890+
self.owner, uuid
891+
)
892+
893+
@async_client_handler(check_no_op=True, check_offline=True)
894+
async def approve_runs(self, uuids: Union[List[str], V1Uuids]):
895+
if isinstance(uuids, list):
896+
uuids = V1Uuids(uuids=uuids)
897+
if self.team:
898+
return await self.client.teams_v1.approve_team_runs(
899+
self.owner, self.team, uuids
900+
)
901+
return await self.client.organizations_v1.approve_organization_runs(
902+
self.owner, body=uuids
903+
)
904+
905+
@async_client_handler(check_no_op=True, check_offline=True)
906+
async def archive_runs(self, uuids: Union[List[str], V1Uuids]):
907+
if isinstance(uuids, list):
908+
uuids = V1Uuids(uuids=uuids)
909+
if self.team:
910+
return await self.client.teams_v1.archive_team_runs(
911+
self.owner, self.team, uuids
912+
)
913+
return await self.client.organizations_v1.archive_organization_runs(
914+
self.owner, body=uuids
915+
)
916+
917+
@async_client_handler(check_no_op=True, check_offline=True)
918+
async def restore_runs(self, uuids: Union[List[str], V1Uuids]):
919+
if isinstance(uuids, list):
920+
uuids = V1Uuids(uuids=uuids)
921+
if self.team:
922+
return await self.client.teams_v1.restore_team_runs(
923+
self.owner, self.team, uuids
924+
)
925+
return await self.client.organizations_v1.restore_organization_runs(
926+
self.owner, body=uuids
927+
)
928+
929+
@async_client_handler(check_no_op=True, check_offline=True)
930+
async def delete_runs(self, uuids: Union[List[str], V1Uuids]):
931+
if isinstance(uuids, list):
932+
uuids = V1Uuids(uuids=uuids)
933+
logger.info("Deleting {} runs".format(len(uuids.uuids)))
934+
if self.team:
935+
return await self.client.teams_v1.delete_team_runs(
936+
self.owner, self.team, body=uuids
937+
)
938+
return await self.client.organizations_v1.delete_organization_runs(
939+
self.owner, body=uuids
940+
)
941+
942+
@async_client_handler(check_no_op=True, check_offline=True)
943+
async def stop_runs(self, uuids: Union[List[str], V1Uuids]):
944+
if isinstance(uuids, list):
945+
uuids = V1Uuids(uuids=uuids)
946+
if self.team:
947+
return await self.client.teams_v1.stop_team_runs(
948+
self.owner, self.team, body=uuids
949+
)
950+
return await self.client.organizations_v1.stop_organization_runs(
951+
self.owner, body=uuids
952+
)
953+
954+
@async_client_handler(check_no_op=True, check_offline=True)
955+
async def skip_runs(self, uuids: Union[List[str], V1Uuids]):
956+
if isinstance(uuids, list):
957+
uuids = V1Uuids(uuids=uuids)
958+
if self.team:
959+
return await self.client.teams_v1.skip_team_runs(
960+
self.owner, self.team, body=uuids
961+
)
962+
return await self.client.organizations_v1.skip_organization_runs(
963+
self.owner, body=uuids
964+
)
965+
966+
@async_client_handler(check_no_op=True, check_offline=True)
967+
async def invalidate_runs(self, uuids: Union[List[str], V1Uuids]):
968+
if isinstance(uuids, list):
969+
uuids = V1Uuids(uuids=uuids)
970+
if self.team:
971+
return await self.client.teams_v1.invalidate_team_runs(
972+
self.owner, self.team, body=uuids
973+
)
974+
return await self.client.organizations_v1.invalidate_organization_runs(
975+
self.owner, body=uuids
976+
)
977+
978+
@async_client_handler(check_no_op=True, check_offline=True)
979+
async def bookmark_runs(self, uuids: Union[List[str], V1Uuids]):
980+
if isinstance(uuids, list):
981+
uuids = V1Uuids(uuids=uuids)
982+
if self.team:
983+
return await self.client.teams_v1.bookmark_team_runs(
984+
self.owner, self.team, body=uuids
985+
)
986+
return await self.client.organizations_v1.bookmark_organization_runs(
987+
self.owner, body=uuids
988+
)
989+
990+
@async_client_handler(check_no_op=True, check_offline=True)
991+
async def tag_runs(self, uuids: Union[List[str], V1Uuids], tags: List[str]):
992+
if isinstance(uuids, list):
993+
uuids = V1Uuids(uuids=uuids)
994+
data = V1EntitiesTags(
995+
uuids=uuids.uuids,
996+
tags=tags,
997+
)
998+
if self.team:
999+
return await self.client.teams_v1.tag_team_runs(
1000+
self.owner, self.team, body=data
1001+
)
1002+
return await self.client.organizations_v1.tag_organization_runs(
1003+
self.owner, body=data
1004+
)
1005+
1006+
@async_client_handler(check_no_op=True, check_offline=True)
1007+
async def transfer_runs(
1008+
self,
1009+
uuids: Union[List[str], V1Uuids],
1010+
to_project: str,
1011+
):
1012+
if isinstance(uuids, list):
1013+
transfer_data = V1EntitiesTransfer(uuids=uuids, project=to_project)
1014+
else:
1015+
transfer_data = V1EntitiesTransfer(uuids=uuids.uuids, project=to_project)
1016+
1017+
logger.info(
1018+
"Transferring {} runs to project {}".format(
1019+
len(transfer_data.uuids), to_project
1020+
)
1021+
)
1022+
if self.team:
1023+
return await self.client.teams_v1.transfer_team_runs(
1024+
self.owner, self.team, body=transfer_data
1025+
)
1026+
return await self.client.organizations_v1.transfer_organization_runs(
1027+
self.owner, body=transfer_data
1028+
)
1029+
1030+
@async_client_handler(check_no_op=True, check_offline=True)
1031+
async def list_versions(
1032+
self,
1033+
kind: V1ProjectVersionKind,
1034+
query: Optional[str] = None,
1035+
sort: Optional[str] = None,
1036+
limit: Optional[int] = None,
1037+
offset: Optional[int] = None,
1038+
) -> V1ListProjectVersionsResponse:
1039+
self._validate_kind(kind)
1040+
params = get_query_params(
1041+
limit=limit or 20, offset=offset, query=query, sort=sort
1042+
)
1043+
if self.team:
1044+
return await self.client.teams_v1.get_team_versions(
1045+
self.owner, self.team, kind, **params
1046+
)
1047+
return await self.client.organizations_v1.get_organization_versions(
1048+
self.owner, kind, **params
1049+
)
1050+
1051+
@async_client_handler(check_no_op=True, check_offline=True)
1052+
async def list_component_versions(
1053+
self,
1054+
query: Optional[str] = None,
1055+
sort: Optional[str] = None,
1056+
limit: Optional[int] = None,
1057+
offset: Optional[int] = None,
1058+
) -> V1ListProjectVersionsResponse:
1059+
return await self.list_versions(
1060+
kind=V1ProjectVersionKind.COMPONENT,
1061+
query=query,
1062+
sort=sort,
1063+
limit=limit,
1064+
offset=offset,
1065+
)
1066+
1067+
@async_client_handler(check_no_op=True, check_offline=True)
1068+
async def list_model_versions(
1069+
self,
1070+
query: Optional[str] = None,
1071+
sort: Optional[str] = None,
1072+
limit: Optional[int] = None,
1073+
offset: Optional[int] = None,
1074+
) -> V1ListProjectVersionsResponse:
1075+
return await self.list_versions(
1076+
kind=V1ProjectVersionKind.MODEL,
1077+
query=query,
1078+
sort=sort,
1079+
limit=limit,
1080+
offset=offset,
1081+
)
1082+
1083+
@async_client_handler(check_no_op=True, check_offline=True)
1084+
async def list_artifact_versions(
1085+
self,
1086+
query: Optional[str] = None,
1087+
sort: Optional[str] = None,
1088+
limit: Optional[int] = None,
1089+
offset: Optional[int] = None,
1090+
) -> V1ListProjectVersionsResponse:
1091+
return await self.list_versions(
1092+
kind=V1ProjectVersionKind.ARTIFACT,
1093+
query=query,
1094+
sort=sort,
1095+
limit=limit,
1096+
offset=offset,
1097+
)
1098+
1099+
@async_client_handler(check_no_op=True, check_offline=True)
1100+
async def list_runs_artifacts_lineage(
1101+
self,
1102+
name: Optional[str] = None,
1103+
query: Optional[str] = None,
1104+
sort: Optional[str] = None,
1105+
limit: Optional[int] = None,
1106+
offset: Optional[int] = None,
1107+
bookmarks: bool = False,
1108+
mode: Optional[str] = None,
1109+
) -> V1ListRunArtifactsResponse:
1110+
params = get_query_params(
1111+
limit=limit or 20, offset=offset, query=query, sort=sort
1112+
)
1113+
if name:
1114+
params["name"] = name
1115+
if bookmarks:
1116+
params["bookmarks"] = bookmarks
1117+
if mode:
1118+
params["mode"] = mode
1119+
1120+
if self.team:
1121+
return await self.client.teams_v1.get_team_runs_artifacts_lineage(
1122+
self.owner, self.team, **params
1123+
)
1124+
return await self.client.organizations_v1.get_organization_runs_artifacts_lineage(
1125+
self.owner, **params
1126+
)

0 commit comments

Comments
 (0)