66
77import  gc 
88import  logging 
9- import  multiprocessing  as  mp 
10- import  os 
11- import  signal 
12- import  subprocess 
139from  abc  import  ABC , abstractmethod 
14- from  datetime  import  datetime 
1510from  time  import  sleep , time 
1611from  typing  import  Any 
1712
1813import  httpx 
19- import  psutil 
2014import  torch 
2115from  fastapi  import  FastAPI 
2216from  langchain .llms .base  import  LLM 
@@ -39,45 +33,21 @@ def load(self) -> Any:
3933	def  offload (self ):
4034		...
4135
42- pid  =  mp .Value ('i' , 0 )
4336
4437class  EmbeddingModelLoader (Loader ):
4538	def  __init__ (self , config : TConfig ):
4639		self .config  =  config 
47- 		logfile_path  =  os .path .join (
48- 			os .environ ['EM_SERVER_LOG_PATH' ],
49- 			f'embedding_server_{ datetime .now ().strftime ("%Y-%m-%d" )}  ,
50- 		)
51- 		self .logfile  =  open (logfile_path , 'a+' )
5240
5341	def  load (self ):
54- 		global  pid 
55- 
5642		emconf  =  self .config .embedding 
57- 
58- 		# start the embedding server if workers are > 0 
59- 		if  emconf .workers  >  0 :
60- 			# compare with None, as PID can be 0, you never know 
61- 			if  pid .value  >  0  and  psutil .pid_exists (pid .value ):
62- 				return 
63- 
64- 			proc  =  subprocess .Popen (  # noqa: S603 
65- 				['./main_em.py' ],
66- 				stdout = self .logfile ,
67- 				stderr = self .logfile ,
68- 				stdin = None ,
69- 				close_fds = True ,
70- 				env = os .environ ,
71- 			)
72- 			pid .value  =  proc .pid 
73- 
7443		last_resp , last_exc  =  None , None 
7544		# poll for heartbeat 
7645		try_  =  0 
77- 		while   try_   <   20 :
78- 			with   httpx . Client ()  as   client :
46+ 		with   httpx . Client ()  as   client :
47+ 			while   try_   <   20 :
7948				try :
8049					# test the server is up 
50+ 					# todo: replace with a tcp connection check 
8151					response  =  client .post (
8252						f'{ emconf .protocol } { emconf .host } { emconf .port }  ,
8353						json = {'input' : 'hello' },
@@ -98,10 +68,7 @@ def load(self):
9868		) from  last_exc 
9969
10070	def  offload (self ):
101- 		global  pid 
102- 		if  pid .value  >  0  and  psutil .pid_exists (pid .value ):
103- 			os .kill (pid .value , signal .SIGTERM )
104- 		self .logfile .close ()
71+ 		...
10572
10673
10774class  VectorDBLoader (Loader ):
0 commit comments