session.py 6.29 KB
Newer Older
Marco Govoni's avatar
Marco Govoni committed
1
2
from __future__ import print_function

3
4
class Session(object):
    """Class for setting up a session, connected to a remove server via rest APIs.
Marco Govoni's avatar
Marco Govoni committed
5
6
7
8
    
    :Example:
    
    >>> from westpy import *
9
    >>> session = Session("email@domain.com")
Marco Govoni's avatar
Marco Govoni committed
10
11
12
    
    """
    def __init__(self,emailId) :
13
        self.token = None
Marco Govoni's avatar
Marco Govoni committed
14
        self.emailId = str(emailId)
15
16
17
        #
        # --- CONFIGURABLE PARAMETERS --- 
        self.serverName = "imedevel.uchicago.edu"
Marco Govoni's avatar
Marco Govoni committed
18
19
20
        self.restAPIinit = "http://imedevel.uchicago.edu:8000/getSessionId"
        self.restAPIrun  = "http://imedevel.uchicago.edu:8000/runWestCode"
        self.restAPIstop = "http://imedevel.uchicago.edu:8000/stopSession"
21
22
23
24
        self.maxTime = 1800 # seconds
        self.maxNumberOfCores = 4
        self.allowedExecutables = ["pw","wstat","wfreq"]
        # -------------------------------
Marco Govoni's avatar
Marco Govoni committed
25
        #
26
        data = {'emailId': self.emailId ,'sessionTime':str(self.maxTime)} 
Marco Govoni's avatar
Marco Govoni committed
27
28
29
30
31
32
33
34
35
        #
        import requests
        import json
        #
        response = None
        try:
            output = requests.post(self.restAPIinit, data=json.dumps(data))
            response = json.loads(output.text)
        except Exception as e:
36
            print('The server is not responding.',e)
Marco Govoni's avatar
Marco Govoni committed
37
38
        if response:
            if "Error" in response:
39
                print("Server failed with the following error \n",response['Error'])
Marco Govoni's avatar
Marco Govoni committed
40
41
42
                return None
            else:
                print("Check the inbox/spam folder of your email and click on the link to activate the session")
43
                self.token = response["token"]
Marco Govoni's avatar
Marco Govoni committed
44
        else:
45
            print('The server is not responding.')
Marco Govoni's avatar
Marco Govoni committed
46
    
47
48
    def getToken(self):
        """Returns the token of the session.
Marco Govoni's avatar
Marco Govoni committed
49
50
51
52
        
        :Example:
        
        >>> from westpy import *
53
54
        >>> session = Session("email@domain.com")
        >>> token = session.getToken()
Marco Govoni's avatar
Marco Govoni committed
55
56
        
        """
57
58
        if self.token :
            return self.token
Marco Govoni's avatar
Marco Govoni committed
59
60
61
62
        else:
            raise ValueError("Cannot find output.")
    
    def stop(self):
63
        """Stops the session and clears the remote workspace. 
Marco Govoni's avatar
Marco Govoni committed
64
65
66
67
        
        :Example:
        
        >>> from westpy import *
68
69
        >>> session = Session("email@domain.com")
        >>> session.stop()
Marco Govoni's avatar
Marco Govoni committed
70
71
72
73
74
75
        
        """
        
        import requests
        import json
        #
76
        headers = {'Content-Type':'application/json; charset=utf-8','emailId':self.emailId,'token':self.token}
Marco Govoni's avatar
Marco Govoni committed
77
78
79
        try:
            response = requests.get(self.restAPIstop, headers=headers, timeout=None)
        except Exception as e:
80
            print('The server is not responding.',e)		 
Marco Govoni's avatar
Marco Govoni committed
81
82
        return json.loads(response.text) 
    
Marco Govoni's avatar
Marco Govoni committed
83
    def run(self,executable=None,inputFile=None,outputFile=None,downloadUrl=[],number_of_cores=2) :
84
        """Runs the executable on the remote server.
Marco Govoni's avatar
Marco Govoni committed
85
86
87
88
89
90
91
        
        :param executable: name of executable
        :type executable: string
        :param inputFile: name of input file
        :type inputFile: string
        :param outputFile: name of output file
        :type outputFile: string
Marco Govoni's avatar
Marco Govoni committed
92
93
        :param downloadUrl: URLs to be downloaded
        :type downloadUrl: list of string
Marco Govoni's avatar
Marco Govoni committed
94
95
96
97
98
99
        :param number_of_cores: number of cores
        :type number_of_cores: int
        
        :Example:
        
        >>> from westpy import *
100
101
102
        >>> session = Session("email@domain.com")
        >>> session.run( "pw", "pw.in", "pw.out", ["http://www.quantum-simulation.org/potentials/sg15_oncv/upf/C_ONCV_PBE-1.0.upf"] , 2 )
        >>> session.stop()
Marco Govoni's avatar
Marco Govoni committed
103
104
105
106
107
        
        """
        #
        import json
        #
108
109
110
111
        assert( number_of_cores <= self.maxNumberOfCores )
        #
        output_dict = {}
        if executable in self.allowedExecutables : 
Marco Govoni's avatar
Marco Govoni committed
112
113
114
115
116
117
           # set inputs
           if inputFile is None:
              inputFile = str(executable)+".in"
           if outputFile is None:
              outputFile = str(executable)+".out"		 
           try:
Marco Govoni's avatar
Marco Govoni committed
118
              output = self.__runExecutable(executable,inputFile,downloadUrl,number_of_cores)
Marco Govoni's avatar
Marco Govoni committed
119
120
              output_json = json.loads(output)
              if "Error" in output_json:
121
                 print("Server failed with the following error \n",output_json['Error'])
Marco Govoni's avatar
Marco Govoni committed
122
123
124
125
126
                 return None			   
              elif "JOB DONE." not in str(output_json['output']).strip():
                 print("MPI execution failed with the following error:  \n"+str(output))
                 return None
              output_data = str(output_json['output']).strip()
127
              output_dict = json.loads(output_json['output_dict'])
Marco Govoni's avatar
Marco Govoni committed
128
129
130
131
132
133
134
              # write the output file
              with open(outputFile, "w") as file :
                 file.write(str(output_data))
           except Exception as e:
              print("Session Expired! Invalid Request sent, Please recreate session and recheck your input. \n"+ e)		 
              return None          
        else:
135
           raise ValueError("Invalid executable name") 
Marco Govoni's avatar
Marco Govoni committed
136
137
138
        #
        print("Generated ",outputFile)
        return output_dict
Marco Govoni's avatar
Marco Govoni committed
139
      
Marco Govoni's avatar
Marco Govoni committed
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
    def __runExecutable(self,executable,input_file,download_urls,number_of_cores) :
        """Runs remotely the executable using a REST api.
        """
        #
        import requests
        import json
        # suck in the input file
        try:
           file_content = ""	     
           with open(input_file,'r') as f :
              for line in f :
                 file_content = file_content + line + "\\n"
        except FileNotFoundError:
           error = "Could not find "+ input_file + ". \n Generate input file "+ input_file +" and try again."	  
           print(error)
           return None
        body = {'urls':download_urls,'file':file_content,'cmd_timeout':'600','script_type':str(executable),'no_of_cores':str(number_of_cores)}  
        jsondata = json.dumps(body)
        jsondataasbytes = jsondata.encode('utf-8')   # needs to be bytes
159
        headers = {'Content-Type':'application/json; charset=utf-8','emailId':self.emailId,'token':self.token}
Marco Govoni's avatar
Marco Govoni committed
160
        response = ""
161
        print("Requested to run executable", executable, "on server", self.serverName)
Marco Govoni's avatar
Marco Govoni committed
162
163
164
165
        try:
           response = requests.post(self.restAPIrun, data = jsondataasbytes, headers=headers, timeout=None)
           return response.text
        except Exception as e:
166
           print('The server is not responding.',e)	
Marco Govoni's avatar
Marco Govoni committed
167
168
           return None
        return None