Coverage for source/paperspace/gradient_handler.py: 81%

1# paperspace/gradient_handler.py 

2 

3from gradient import NotebooksClient 

4import os 

5from datetime import datetime 

6 

7""" 

8Deafult start command for Paperspace Gradient Notebooks. Needed to be passed in order for Notebook to be fully functional. 

9""" 

10DEFAULT_START_COMMAND = " & PIP_DISABLE_PIP_VERSION_CHECK=1 jupyter lab --allow-root --ip=0.0.0.0 --no-browser \ 

11 --ServerApp.trust_xheaders=True --ServerApp.disable_check_xsrf=False --ServerApp.allow_remote_access=True \ 

12 --ServerApp.allow_origin='*' --ServerApp.allow_credentials=True" 

13 

14""" 

15Default container type for Paperspace Gradient Notebooks. It is a base container with preinstalled Tensorflow, JAX and Python 3.9. 

16""" 

17DEFAULT_CONTAINER_TYPE = 'paperspace/gradient-base:pt112-tf29-jax0317-py39-20230125' 

18 

19class GradientHandler(): 

20 """ 

21 Responsible for communication and management of Paperspace Gradient services. 

22 """ 

23 

24 def __init__(self) -> None: 

25 """ 

26 Class constructor. Before calling it GRADIENT_API_KEY and GRADIENT_PROJECT_ID 

27 should be available in as environmental variables. 

28 

29 Raises: 

30 RuntimeError: If Paperspace Gradient api key or project ID are not defined. 

31 """ 

32 

33 GRADIENT_API_KEY = os.getenv('GRADIENT_API_KEY') 

34 GRADIENT_PROJECT_ID = os.getenv('GRADIENT_PROJECT_ID') 

35 

36 if not GRADIENT_API_KEY or not GRADIENT_PROJECT_ID: 

37 raise RuntimeError('Paperspace Gradient api key or project ID not found in environment variables!') 

38 

39 self.notebooks = NotebooksClient(GRADIENT_API_KEY) 

40 self.project_id = GRADIENT_PROJECT_ID 

41 

42 def create_notebook(self, command_to_invoke: str, github_repository_url: str = None, 

43 notebook_name: str = datetime.now(), machine_types: list = ['Free-P5000'], 

44 timeout: int = 6, environment_dict: dict = dict()) -> str: 

45 """ 

46 Attempts to create notebook basing on certain github repository, starting command and 

47 sets needed environmental parameters (e.g. variables, machine type, etc.). 

48 

49 Parameters: 

50 command_to_invoke (str): Command to be invoked after notebook is started. 

51 github_repository_url (str): URL to repository that should be downloaded to notebook. 

52 notebook_name (str): Name that should be given to notebook. 

53 machine_types (list): List of demanded machine's types that notebook should be attempted 

54 to be created on. The first successful creation stops attempts to create notebook for 

55 machine types that are further in the list. Therefore, machine's types should be ordered 

56 from the most wanted to the least one. 

57 timeout (int): Number of hours that notebook should be active for. For free instances maximal 

58 value for this parameter is 6. 

59 environment_dict (dict): Dictionary containing defined environmental variables. 

60 

61 Raises: 

62 RuntimeError: If approached problem during notebook creation. 

63 """ 

64 

65 notebook_id = None 

66 error_to_be_raised = None 

67 for machine_type in machine_types: 

68 try: 

69 notebook_id = self.notebooks.create(machine_type = machine_type, 

70 container = DEFAULT_CONTAINER_TYPE, 

71 project_id = self.project_id, 

72 shutdown_timeout = timeout, 

73 workspace = github_repository_url, 

74 command = command_to_invoke + DEFAULT_START_COMMAND, 

75 environment = environment_dict, 

76 name = notebook_name) 

77 except Exception as error: 

78 error_to_be_raised = error 

79 

80 if notebook_id is not None: 

81 break 

82 

83 if notebook_id is None: 

84 raise RuntimeError(f"Did not managed to create notebook! Original error: {error_to_be_raised}") 

85 

86 return notebook_id 

87 

88 def delete_notebook(self, notebook_id: str) -> None: 

89 """ 

90 Deletes notebook with the given ID from Paperspace Gradient. 

91 

92 This method attempts to remove a previously created notebook from 

93 the Paperspace Gradient platform. 

94 

95 Parameters: 

96 notebook_id (str): ID of the notebook to be deleted. 

97 

98 Raises: 

99 RuntimeError: If approached problem during notebook deletion. 

100 """ 

101 

102 try: 

103 self.notebooks.delete(notebook_id) 

104 except Exception as error: 

105 raise RuntimeError(f"Did not managed to delete notebook! Original error: {error}")