Coverage for melissa/server/main.py: 54%

1"""This script is the main Melissa server entrypoint.""" 

2 

3import argparse 

4import importlib.util 

5import sys 

6import logging 

7import os 

8from pathlib import Path 

9from typing import Any, Dict 

10 

11import rapidjson 

12 

13from melissa.launcher.__main__ import validate_config, CONFIG_PARSE_MODE 

14from melissa.server.base_server import BaseServer 

15from melissa.server.exceptions import MelissaError 

16 

17logger = logging.getLogger(__name__) 

18 

19 

20def main() -> None: 

21 """ 

22 Initializes and starts the Melissa server. 

23 

24 ### Overview 

25 This function is the entry point for the Melissa server. It is **not intended to be called 

26 directly by users** but is invoked by the `melissa-launcher`. 

27 

28 ### Steps Performed 

29 1. Parses command-line arguments, including flags like `--config_name` to 

30 locate the project directory. 

31 2. Resolves the appropriate custom server class, if specified in the configuration. 

32 3. Instantiates the server using the user-provided configuration. 

33 4. Starts the server by calling `server.start()`. 

34 

35 ### Usage 

36 This function is typically invoked through the server shell script `server.sh`: 

37 

38 ``` 

39 exec python -m melissa --config_name /path/to/config 

40 ```""" 

41 

42 parser = argparse.ArgumentParser( 

43 prog="melissa", description="Melissa" 

44 ) 

45 

46 # CLI flags 

47 parser.add_argument( 

48 "--project_dir", 

49 help="Directory to all necessary files:\n" 

50 " client.sh\n" 

51 " server.sh\n" 

52 " config.py\n" 

53 " data_generator\n" 

54 " CustomServer.py" 

55 ) 

56 

57 parser.add_argument( 

58 "--config_name", 

59 help="Defaults to `config` but user can change" 

60 "the search name by indicating it with this flag", 

61 default=None 

62 ) 

63 

64 args = parser.parse_args() 

65 if not args.config_name: 

66 conf_name = 'config' 

67 else: 

68 conf_name = args.config_name 

69 

70 # load the config into a python dictionary 

71 with open(Path(args.project_dir) / f"{conf_name}.json") as json_file: 

72 config_dict = rapidjson.load(json_file, parse_mode=CONFIG_PARSE_MODE) 

73 

74 config_dict['user_data_dir'] = Path('user_data') 

75 

76 # ensure user passed the necessary information for software configuration 

77 args, config_dict = validate_config(args, config_dict) 

78 

79 # Resolve the server 

80 myserver = get_resolved_server(args, config_dict) 

81 

82 try: 

83 myserver.configure_logger() 

84 myserver.initialize_connections() 

85 myserver.start() 

86 except MelissaError as e: 

87 logger.exception(e) 

88 if myserver.no_fault_tolerance: 

89 myserver.close_connection(1) 

90 myserver.mpi_abort(1) 

91 

92 

93def safe_join(base_path: str, new_path: str) -> Path: 

94 """Safely joins a base path and a new path, ensuring the result is an expanded absolute path. 

95 

96 ### Parameters 

97 - **base_path (str)**: The base directory path. 

98 - **new_path (str)**: The path to join with the base path. 

99 

100 ### Returns 

101 - **Path**: An expanded and absolute path combining `base_path` and `new_path`.""" 

102 

103 new_path = os.path.expandvars(new_path) 

104 new_path_ = Path(new_path) 

105 if new_path_.is_absolute(): 

106 return new_path_ 

107 return Path(base_path) / new_path_ 

108 

109 

110def get_resolved_server(args: argparse.Namespace, 

111 config_dict: Dict[str, Any]) -> BaseServer: 

112 """Resolves and returns a server object based on the provided configurations. 

113 

114 ### Parameters 

115 - **args (argparse.Namespace)**: Parsed command-line arguments containing 

116 server-related options. 

117 - **config_dict (Dict[str, Any])**: Configuration dictionary with server settings. 

118 

119 ### Returns 

120 - `BaseServer`: An instance of the resolved server class, initialized with the 

121 provided configurations. 

122 

123 ### Raises 

124 - `ImportError`: If the server file or class cannot be imported. 

125 - `AttributeError`: If the specified server class is not found in the module. 

126 - `Exception`: If there is an issue with initializing the server class.""" 

127 

128 # Resolve the server file name and path 

129 server_file_name = config_dict.get("server_filename", "server.py") 

130 server_path = safe_join(args.project_dir, server_file_name) 

131 server_module_name = Path(server_path).stem 

132 # Import the server module dynamically 

133 spec_server = importlib.util.spec_from_file_location(server_module_name, server_path) 

134 if not (spec_server and spec_server.loader): 

135 logger.warning("Unable to import server module from path.") 

136 raise ImportError(f"Failed to import server module from {server_path}.") 

137 

138 sys.path.append(str(Path(server_path).parent)) 

139 server = importlib.util.module_from_spec(spec_server) 

140 spec_server.loader.exec_module(server) 

141 

142 # Resolve the server class name 

143 server_class_name = config_dict.get("server_class", "MyServer") 

144 try: 

145 MyServerClass = getattr(server, server_class_name) 

146 except AttributeError: 

147 raise AttributeError( 

148 f"Server class '{server_class_name}' not found in '{server_file_name}'." 

149 ) 

150 

151 return MyServerClass(config_dict)