wallaroo.workspace
1import json 2from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast 3 4import requests 5 6from wallaroo.wallaroo_ml_ops_api_client.models.workspaces_create_response_200 import ( 7 WorkspacesCreateResponse200, 8) 9 10from . import queries 11from .models import Models, ModelsList 12from .object import * 13from .pipeline import Pipeline 14from .user import User 15from .user_type import UserType 16from .version import _user_agent 17from .wallaroo_ml_ops_api_client.api.workspace import workspaces_create 18from .wallaroo_ml_ops_api_client.models import workspaces_create_json_body 19 20if TYPE_CHECKING: 21 # Imports that happen below in methods to fix circular import dependency 22 # issues need to also be specified here to satisfy mypy type checking. 23 from .client import Client 24 25 26class Workspace(Object): 27 """Workspace provides a user and visibility context for access to models and pipelines.""" 28 29 def __init__( 30 self, client: Optional["Client"], data: Dict[str, Any], standalone=False 31 ) -> None: 32 assert client is not None 33 self.client = client 34 # data = {"name": workspace_name, "owner": owner_user_id} 35 super().__init__(gql_client=client._gql_client, data=data) 36 37 def __repr__(self) -> str: 38 return str(self.to_json()) 39 40 def _fetch_attributes(self) -> Dict[str, Any]: 41 """Fetches all member data from the GraphQL API.""" 42 data = self._gql_client.execute( 43 gql.gql(queries.named("WorkspaceById")), 44 variable_values={"id": self.id()}, 45 )["workspace_by_pk"] 46 47 return data 48 49 def _fill(self, data: Dict[str, Any]) -> None: 50 """Fills an object given a response dictionary from the GraphQL API. 51 52 Only the primary key member must be present; other members will be 53 filled in via rehydration if their corresponding member function is 54 called. 55 """ 56 57 for required_attribute in ["id"]: 58 if required_attribute not in data: 59 raise RequiredAttributeMissing( 60 self.__class__.__name__, required_attribute 61 ) 62 # Required 63 self._id = data["id"] 64 65 # Optional 66 self._name = value_if_present(data, "name") 67 self._archived = value_if_present(data, "archived") 68 self._created_at = value_if_present(data, "created_at") 69 self._created_by = value_if_present(data, "created_by") 70 self._models = value_if_present(data, "models") 71 self._pipelines = value_if_present(data, "pipelines") 72 self._users = value_if_present(data, "users") 73 74 def to_json(self): 75 return { 76 "name": self.name(), 77 "id": self.id(), 78 "archived": self.archived(), 79 "created_by": self.created_by(), 80 "created_at": self.created_at(), 81 "models": self.models(), 82 "pipelines": self.pipelines(), 83 } 84 85 @staticmethod 86 def _get_user_default_workspace(client): 87 88 res = client._gql_client.execute( 89 gql.gql(queries.named("UserDefaultWorkspace")), 90 variable_values={"user_id": client.auth.user_id()}, 91 ) 92 if res["user_default_workspace"] == []: 93 return None 94 else: 95 return Workspace(client, res["user_default_workspace"][0]["workspace"]) 96 97 @staticmethod 98 def _create_user_default_workspace(client) -> "Workspace": 99 100 res = client._gql_client.execute( 101 gql.gql(queries.named("CreateDefaultUserWorkspace")), 102 variable_values={ 103 "user_id": client.auth.user_id(), 104 "workspace_name": f"{client.auth.user_email()} - Default Workspace", 105 "user_type": UserType.OWNER.name, 106 }, 107 ) 108 return Workspace(client, res["insert_workspace_one"]) 109 110 @staticmethod 111 def _create_workspace(client: "Client", name: str): 112 body = workspaces_create_json_body.WorkspacesCreateJsonBody(name) 113 res = workspaces_create.sync(client=client.mlops(), json_body=body) 114 115 if not isinstance(res, WorkspacesCreateResponse200): 116 raise Exception("Failed to create workspace.") 117 118 if res is None: 119 raise Exception("Failed to create workspace.") 120 121 return Workspace(client, {"id": res.workspace_id}) 122 123 def id(self) -> int: 124 return self._id 125 126 @rehydrate("_name") 127 def name(self) -> str: 128 return cast(str, self._name) 129 130 @rehydrate("_archived") 131 def archived(self) -> bool: 132 return cast(bool, self._archived) 133 134 @rehydrate("_created_at") 135 def created_at(self) -> str: 136 return cast(str, self._created_at) 137 138 @rehydrate("_created_by") 139 def created_by(self) -> str: 140 return cast(str, self._created_by) 141 142 @rehydrate("_models") 143 def models(self) -> ModelsList: 144 """ 145 Returns a List of Models objects that have been created in this workspace. 146 """ 147 self._fill(self._fetch_attributes()) 148 return cast( 149 ModelsList, 150 list( 151 map( 152 lambda m: Models(self.client, data=m), 153 cast(List[Dict], self._models), 154 ) 155 ), 156 ) 157 158 @rehydrate("_pipelines") 159 def pipelines(self) -> List[Pipeline]: 160 self._fill(self._fetch_attributes()) 161 return cast( 162 List[Pipeline], 163 list( 164 map( 165 lambda m: Pipeline(self.client, data=m), 166 cast(List[Dict], self._pipelines), 167 ) 168 ), 169 ) 170 171 @rehydrate("_users") 172 def users(self) -> List[User]: 173 return cast( 174 List[User], 175 list( 176 map( 177 lambda u: self.client._get_user_by_id(u["user_id"]), 178 cast(List[Dict], self._users), 179 ) 180 ), 181 ) 182 183 def add_user(self, user_email: str) -> "Workspace": 184 """Add a user to workspace as participant""" 185 user = self.client.get_user_by_email(user_email) 186 if user is None: 187 raise EntityNotFoundError("User", {"email": user_email}) 188 return self._add_user(user, UserType.COLLABORATOR) 189 190 def add_owner(self, user_email: str) -> "Workspace": 191 """Add a user to workspace as owner""" 192 193 user = self.client.get_user_by_email(user_email) 194 if user is None: 195 raise EntityNotFoundError("User", {"email": user_email}) 196 197 return self._add_user(user, UserType.OWNER) 198 199 def _add_user(self, user: User, user_type: UserType) -> "Workspace": 200 headers = { 201 "authorization": self.client.auth._bearer_token_str(), 202 "user-agent": _user_agent, 203 } 204 invite = requests.post( 205 f"{self.client.api_endpoint}/v1/api/workspaces/add_user", 206 data=json.dumps( 207 { 208 "email": user.email(), 209 "workspace_id": self._id, 210 "user_id": user.id(), 211 "url": "http://localhost:3000", 212 "user_type": user_type, 213 } 214 ), 215 headers=headers, 216 ) 217 if invite.status_code > 299: 218 raise Exception(f"Failed to invite: {user.id()} to workspace: {self._id}") 219 self._fill(data={"id": self._id}) 220 return self 221 222 def remove_user(self, user_email: str): 223 224 user = self.client.get_user_by_email(user_email) 225 assert user is not None 226 if user.id() == self.client.auth.user_id(): 227 default_workspace = self._get_user_default_workspace(self.client) 228 if default_workspace.id() == self.id(): 229 raise Exception("User cannot be removed from their default workspace") 230 231 if len(self.users()) <= 1: 232 raise Exception("Last user of a workspace cannot be removed") 233 234 res = self.client._gql_client.execute( 235 gql.gql(queries.named("DeleteWorkspaceUser")), 236 variable_values={"user_id": user.id(), "workspace_id": self.id()}, 237 ) 238 if user.id() == self.client.auth.user_id(): 239 self._fill({"id": default_workspace.id()}) 240 else: 241 self._rehydrate() 242 243 244class Workspaces(List[Workspace]): 245 """Wraps a list of workspaces for display.""" 246 247 @staticmethod 248 def _format_users(users: List[User]) -> List[str]: 249 return [u.email() for u in users] 250 251 def _repr_html_(self) -> str: 252 rows = [ 253 f""" 254 <tr > 255 <td>{r.name()}</td> 256 <td>{r.created_at()[:19].replace("T", " ")}</td> 257 <td>{Workspaces._format_users(r.users())}</td> 258 <td>{len(r.models())}</td> 259 <td>{len(r.pipelines())}</td> 260 </tr> 261 """ 262 for r in self 263 ] 264 table = """ 265 <table> 266 <tr> 267 <th>Name</th> 268 <th>Created At</th> 269 <th>Users</th> 270 <th>Models</th> 271 <th>Pipelines</th> 272 </tr> 273 {0} 274 </table> 275 """.format( 276 "\n".join(rows) 277 ) 278 return table
27class Workspace(Object): 28 """Workspace provides a user and visibility context for access to models and pipelines.""" 29 30 def __init__( 31 self, client: Optional["Client"], data: Dict[str, Any], standalone=False 32 ) -> None: 33 assert client is not None 34 self.client = client 35 # data = {"name": workspace_name, "owner": owner_user_id} 36 super().__init__(gql_client=client._gql_client, data=data) 37 38 def __repr__(self) -> str: 39 return str(self.to_json()) 40 41 def _fetch_attributes(self) -> Dict[str, Any]: 42 """Fetches all member data from the GraphQL API.""" 43 data = self._gql_client.execute( 44 gql.gql(queries.named("WorkspaceById")), 45 variable_values={"id": self.id()}, 46 )["workspace_by_pk"] 47 48 return data 49 50 def _fill(self, data: Dict[str, Any]) -> None: 51 """Fills an object given a response dictionary from the GraphQL API. 52 53 Only the primary key member must be present; other members will be 54 filled in via rehydration if their corresponding member function is 55 called. 56 """ 57 58 for required_attribute in ["id"]: 59 if required_attribute not in data: 60 raise RequiredAttributeMissing( 61 self.__class__.__name__, required_attribute 62 ) 63 # Required 64 self._id = data["id"] 65 66 # Optional 67 self._name = value_if_present(data, "name") 68 self._archived = value_if_present(data, "archived") 69 self._created_at = value_if_present(data, "created_at") 70 self._created_by = value_if_present(data, "created_by") 71 self._models = value_if_present(data, "models") 72 self._pipelines = value_if_present(data, "pipelines") 73 self._users = value_if_present(data, "users") 74 75 def to_json(self): 76 return { 77 "name": self.name(), 78 "id": self.id(), 79 "archived": self.archived(), 80 "created_by": self.created_by(), 81 "created_at": self.created_at(), 82 "models": self.models(), 83 "pipelines": self.pipelines(), 84 } 85 86 @staticmethod 87 def _get_user_default_workspace(client): 88 89 res = client._gql_client.execute( 90 gql.gql(queries.named("UserDefaultWorkspace")), 91 variable_values={"user_id": client.auth.user_id()}, 92 ) 93 if res["user_default_workspace"] == []: 94 return None 95 else: 96 return Workspace(client, res["user_default_workspace"][0]["workspace"]) 97 98 @staticmethod 99 def _create_user_default_workspace(client) -> "Workspace": 100 101 res = client._gql_client.execute( 102 gql.gql(queries.named("CreateDefaultUserWorkspace")), 103 variable_values={ 104 "user_id": client.auth.user_id(), 105 "workspace_name": f"{client.auth.user_email()} - Default Workspace", 106 "user_type": UserType.OWNER.name, 107 }, 108 ) 109 return Workspace(client, res["insert_workspace_one"]) 110 111 @staticmethod 112 def _create_workspace(client: "Client", name: str): 113 body = workspaces_create_json_body.WorkspacesCreateJsonBody(name) 114 res = workspaces_create.sync(client=client.mlops(), json_body=body) 115 116 if not isinstance(res, WorkspacesCreateResponse200): 117 raise Exception("Failed to create workspace.") 118 119 if res is None: 120 raise Exception("Failed to create workspace.") 121 122 return Workspace(client, {"id": res.workspace_id}) 123 124 def id(self) -> int: 125 return self._id 126 127 @rehydrate("_name") 128 def name(self) -> str: 129 return cast(str, self._name) 130 131 @rehydrate("_archived") 132 def archived(self) -> bool: 133 return cast(bool, self._archived) 134 135 @rehydrate("_created_at") 136 def created_at(self) -> str: 137 return cast(str, self._created_at) 138 139 @rehydrate("_created_by") 140 def created_by(self) -> str: 141 return cast(str, self._created_by) 142 143 @rehydrate("_models") 144 def models(self) -> ModelsList: 145 """ 146 Returns a List of Models objects that have been created in this workspace. 147 """ 148 self._fill(self._fetch_attributes()) 149 return cast( 150 ModelsList, 151 list( 152 map( 153 lambda m: Models(self.client, data=m), 154 cast(List[Dict], self._models), 155 ) 156 ), 157 ) 158 159 @rehydrate("_pipelines") 160 def pipelines(self) -> List[Pipeline]: 161 self._fill(self._fetch_attributes()) 162 return cast( 163 List[Pipeline], 164 list( 165 map( 166 lambda m: Pipeline(self.client, data=m), 167 cast(List[Dict], self._pipelines), 168 ) 169 ), 170 ) 171 172 @rehydrate("_users") 173 def users(self) -> List[User]: 174 return cast( 175 List[User], 176 list( 177 map( 178 lambda u: self.client._get_user_by_id(u["user_id"]), 179 cast(List[Dict], self._users), 180 ) 181 ), 182 ) 183 184 def add_user(self, user_email: str) -> "Workspace": 185 """Add a user to workspace as participant""" 186 user = self.client.get_user_by_email(user_email) 187 if user is None: 188 raise EntityNotFoundError("User", {"email": user_email}) 189 return self._add_user(user, UserType.COLLABORATOR) 190 191 def add_owner(self, user_email: str) -> "Workspace": 192 """Add a user to workspace as owner""" 193 194 user = self.client.get_user_by_email(user_email) 195 if user is None: 196 raise EntityNotFoundError("User", {"email": user_email}) 197 198 return self._add_user(user, UserType.OWNER) 199 200 def _add_user(self, user: User, user_type: UserType) -> "Workspace": 201 headers = { 202 "authorization": self.client.auth._bearer_token_str(), 203 "user-agent": _user_agent, 204 } 205 invite = requests.post( 206 f"{self.client.api_endpoint}/v1/api/workspaces/add_user", 207 data=json.dumps( 208 { 209 "email": user.email(), 210 "workspace_id": self._id, 211 "user_id": user.id(), 212 "url": "http://localhost:3000", 213 "user_type": user_type, 214 } 215 ), 216 headers=headers, 217 ) 218 if invite.status_code > 299: 219 raise Exception(f"Failed to invite: {user.id()} to workspace: {self._id}") 220 self._fill(data={"id": self._id}) 221 return self 222 223 def remove_user(self, user_email: str): 224 225 user = self.client.get_user_by_email(user_email) 226 assert user is not None 227 if user.id() == self.client.auth.user_id(): 228 default_workspace = self._get_user_default_workspace(self.client) 229 if default_workspace.id() == self.id(): 230 raise Exception("User cannot be removed from their default workspace") 231 232 if len(self.users()) <= 1: 233 raise Exception("Last user of a workspace cannot be removed") 234 235 res = self.client._gql_client.execute( 236 gql.gql(queries.named("DeleteWorkspaceUser")), 237 variable_values={"user_id": user.id(), "workspace_id": self.id()}, 238 ) 239 if user.id() == self.client.auth.user_id(): 240 self._fill({"id": default_workspace.id()}) 241 else: 242 self._rehydrate()
Workspace provides a user and visibility context for access to models and pipelines.
Workspace( client: Optional[wallaroo.client.Client], data: Dict[str, Any], standalone=False)
30 def __init__( 31 self, client: Optional["Client"], data: Dict[str, Any], standalone=False 32 ) -> None: 33 assert client is not None 34 self.client = client 35 # data = {"name": workspace_name, "owner": owner_user_id} 36 super().__init__(gql_client=client._gql_client, data=data)
Base constructor.
Each object requires:
- a GraphQL client - in order to fill its missing members dynamically
- an initial data blob - typically from unserialized JSON, contains at
- least the data for required members (typically the object's primary key) and optionally other data members.
def
name(*args, **kwargs):
41 def wrapper(*args, **kwargs): 42 obj = args[0] 43 if not getattr(obj, "_standalone", None): 44 present = getattr(obj, attr) != DehydratedValue() 45 # Uncomment to debug while testing 46 # print( 47 # "rehydrate: {} -> {}".format( 48 # attr, "present" if present else "not present" 49 # ) 50 # ) 51 if not present: 52 obj._rehydrate() 53 result = fn(*args, **kwargs) 54 return result
def
archived(*args, **kwargs):
41 def wrapper(*args, **kwargs): 42 obj = args[0] 43 if not getattr(obj, "_standalone", None): 44 present = getattr(obj, attr) != DehydratedValue() 45 # Uncomment to debug while testing 46 # print( 47 # "rehydrate: {} -> {}".format( 48 # attr, "present" if present else "not present" 49 # ) 50 # ) 51 if not present: 52 obj._rehydrate() 53 result = fn(*args, **kwargs) 54 return result
def
created_at(*args, **kwargs):
41 def wrapper(*args, **kwargs): 42 obj = args[0] 43 if not getattr(obj, "_standalone", None): 44 present = getattr(obj, attr) != DehydratedValue() 45 # Uncomment to debug while testing 46 # print( 47 # "rehydrate: {} -> {}".format( 48 # attr, "present" if present else "not present" 49 # ) 50 # ) 51 if not present: 52 obj._rehydrate() 53 result = fn(*args, **kwargs) 54 return result
def
created_by(*args, **kwargs):
41 def wrapper(*args, **kwargs): 42 obj = args[0] 43 if not getattr(obj, "_standalone", None): 44 present = getattr(obj, attr) != DehydratedValue() 45 # Uncomment to debug while testing 46 # print( 47 # "rehydrate: {} -> {}".format( 48 # attr, "present" if present else "not present" 49 # ) 50 # ) 51 if not present: 52 obj._rehydrate() 53 result = fn(*args, **kwargs) 54 return result
def
models(*args, **kwargs):
41 def wrapper(*args, **kwargs): 42 obj = args[0] 43 if not getattr(obj, "_standalone", None): 44 present = getattr(obj, attr) != DehydratedValue() 45 # Uncomment to debug while testing 46 # print( 47 # "rehydrate: {} -> {}".format( 48 # attr, "present" if present else "not present" 49 # ) 50 # ) 51 if not present: 52 obj._rehydrate() 53 result = fn(*args, **kwargs) 54 return result
Returns a List of Models objects that have been created in this workspace.
def
pipelines(*args, **kwargs):
41 def wrapper(*args, **kwargs): 42 obj = args[0] 43 if not getattr(obj, "_standalone", None): 44 present = getattr(obj, attr) != DehydratedValue() 45 # Uncomment to debug while testing 46 # print( 47 # "rehydrate: {} -> {}".format( 48 # attr, "present" if present else "not present" 49 # ) 50 # ) 51 if not present: 52 obj._rehydrate() 53 result = fn(*args, **kwargs) 54 return result
def
users(*args, **kwargs):
41 def wrapper(*args, **kwargs): 42 obj = args[0] 43 if not getattr(obj, "_standalone", None): 44 present = getattr(obj, attr) != DehydratedValue() 45 # Uncomment to debug while testing 46 # print( 47 # "rehydrate: {} -> {}".format( 48 # attr, "present" if present else "not present" 49 # ) 50 # ) 51 if not present: 52 obj._rehydrate() 53 result = fn(*args, **kwargs) 54 return result
184 def add_user(self, user_email: str) -> "Workspace": 185 """Add a user to workspace as participant""" 186 user = self.client.get_user_by_email(user_email) 187 if user is None: 188 raise EntityNotFoundError("User", {"email": user_email}) 189 return self._add_user(user, UserType.COLLABORATOR)
Add a user to workspace as participant
191 def add_owner(self, user_email: str) -> "Workspace": 192 """Add a user to workspace as owner""" 193 194 user = self.client.get_user_by_email(user_email) 195 if user is None: 196 raise EntityNotFoundError("User", {"email": user_email}) 197 198 return self._add_user(user, UserType.OWNER)
Add a user to workspace as owner
def
remove_user(self, user_email: str):
223 def remove_user(self, user_email: str): 224 225 user = self.client.get_user_by_email(user_email) 226 assert user is not None 227 if user.id() == self.client.auth.user_id(): 228 default_workspace = self._get_user_default_workspace(self.client) 229 if default_workspace.id() == self.id(): 230 raise Exception("User cannot be removed from their default workspace") 231 232 if len(self.users()) <= 1: 233 raise Exception("Last user of a workspace cannot be removed") 234 235 res = self.client._gql_client.execute( 236 gql.gql(queries.named("DeleteWorkspaceUser")), 237 variable_values={"user_id": user.id(), "workspace_id": self.id()}, 238 ) 239 if user.id() == self.client.auth.user_id(): 240 self._fill({"id": default_workspace.id()}) 241 else: 242 self._rehydrate()
class
Workspaces(typing.List[wallaroo.workspace.Workspace]):
245class Workspaces(List[Workspace]): 246 """Wraps a list of workspaces for display.""" 247 248 @staticmethod 249 def _format_users(users: List[User]) -> List[str]: 250 return [u.email() for u in users] 251 252 def _repr_html_(self) -> str: 253 rows = [ 254 f""" 255 <tr > 256 <td>{r.name()}</td> 257 <td>{r.created_at()[:19].replace("T", " ")}</td> 258 <td>{Workspaces._format_users(r.users())}</td> 259 <td>{len(r.models())}</td> 260 <td>{len(r.pipelines())}</td> 261 </tr> 262 """ 263 for r in self 264 ] 265 table = """ 266 <table> 267 <tr> 268 <th>Name</th> 269 <th>Created At</th> 270 <th>Users</th> 271 <th>Models</th> 272 <th>Pipelines</th> 273 </tr> 274 {0} 275 </table> 276 """.format( 277 "\n".join(rows) 278 ) 279 return table
Wraps a list of workspaces for display.
Inherited Members
- builtins.list
- list
- clear
- copy
- append
- insert
- extend
- pop
- remove
- index
- count
- reverse
- sort