Skip to content

Commit 3c2f1da

Browse files
UtkarshTheDevUtkarshTheDev
committed
Fixed ngrok post uses and server issues
1 parent dbf55a4 commit 3c2f1da

File tree

4 files changed

+35
-185
lines changed

4 files changed

+35
-185
lines changed

‎locallab/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
LocalLab - A lightweight AI inference server for running LLMs locally
33
"""
44

5-
__version__ = "0.4.38"
5+
__version__ = "0.4.39"
66

77
# Only import what's necessary initially, lazy-load the rest
88
from .logger import get_logger

‎locallab/core/app.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ def init(backend, **kwargs):
3737
DEFAULT_MODEL,
3838
ENABLE_COMPRESSION,
3939
QUANTIZATION_TYPE,
40+
SERVER_PORT,
4041
)
4142
from ..cli.config import get_config_value
4243

@@ -96,7 +97,8 @@ async def startup_event():
9697
use_ngrok = get_config_value("use_ngrok", False)
9798
if use_ngrok:
9899
from ..utils.networking import setup_ngrok
99-
ngrok_url = await setup_ngrok(SERVER_PORT)
100+
port = int(os.environ.get("LOCALLAB_PORT", SERVER_PORT)) # Use SERVER_PORT as fallback
101+
ngrok_url = await setup_ngrok(port)
100102
if ngrok_url:
101103
logger.info(f"{Fore.GREEN}Ngrok tunnel established successfully{Style.RESET_ALL}")
102104

‎locallab/server.py

Lines changed: 30 additions & 182 deletions
Original file line numberDiff line numberDiff line change
@@ -424,6 +424,11 @@ async def serve(self, sock=None):
424424

425425

426426
class ServerWithCallback(uvicorn.Server):
427+
def __init__(self, config):
428+
super().__init__(config)
429+
self.servers = [] # Initialize servers list
430+
self.should_exit = False
431+
427432
def install_signal_handlers(self):
428433
def handle_exit(signum, frame):
429434
self.should_exit = True
@@ -436,207 +441,50 @@ async def startup(self, sockets=None):
436441
if self.should_exit:
437442
return
438443

439-
if sockets is not None:
440-
try:
441-
await super().startup(sockets=sockets)
442-
logger.info("Using uvicorn's built-in Server implementation")
443-
except Exception as e:
444-
logger.error(f"Error during server startup: {str(e)}")
445-
logger.debug(f"Server startup error details: {traceback.format_exc()}")
446-
self.servers = []
444+
try:
445+
await super().startup(sockets=sockets)
446+
logger.info("Using uvicorn's built-in Server implementation")
447+
except Exception as e:
448+
logger.error(f"Error during server startup: {str(e)}")
449+
logger.debug(f"Server startup error details: {traceback.format_exc()}")
450+
self.servers = []
451+
if sockets:
447452
for socket in sockets:
448453
server = SimpleTCPServer(config=self.config)
449454
server.server = self
450455
await server.start()
451456
self.servers.append(server)
452-
else:
453-
try:
454-
await super().startup(sockets=None)
455-
logger.info("Using uvicorn's built-in Server implementation")
456-
except Exception as e:
457-
logger.error(f"Error during server startup: {str(e)}")
458-
logger.debug(f"Server startup error details: {traceback.format_exc()}")
457+
else:
459458
server = SimpleTCPServer(config=self.config)
460459
server.server = self
461460
await server.start()
462-
self.servers = [server]
463-
464-
if self.lifespan is not None:
465-
try:
466-
await self.lifespan.startup()
467-
except Exception as e:
468-
logger.error(f"Error during lifespan startup: {str(e)}")
469-
logger.debug(f"Lifespan startup error details: {traceback.format_exc()}")
470-
self.lifespan = NoopLifespan(self.config.app)
471-
logger.warning("Replaced failed lifespan with NoopLifespan")
472-
473-
async def main_loop(self):
474-
try:
475-
while not self.should_exit:
476-
await asyncio.sleep(0.05)
477-
478-
if self.should_exit:
479-
logger.debug("Shutdown signal detected in main_loop")
480-
break
481-
except Exception as e:
482-
logger.error(f"Error in main loop: {str(e)}")
483-
logger.debug(f"Main loop error details: {traceback.format_exc()}")
484-
self.should_exit = True
485-
486-
logger.debug("Exiting main_loop")
461+
self.servers.append(server)
487462

488463
async def shutdown(self, sockets=None):
489464
logger.debug("Starting server shutdown process")
490465

491-
if self.servers:
492-
for server in self.servers:
493-
try:
494-
server.close()
495-
await server.wait_closed()
496-
logger.debug("Server closed successfully")
497-
except Exception as e:
498-
logger.error(f"Error closing server: {str(e)}")
499-
logger.debug(f"Server close error details: {traceback.format_exc()}")
466+
# First shutdown any SimpleTCPServer instances
467+
for server in self.servers:
468+
try:
469+
await server.shutdown()
470+
logger.debug("SimpleTCPServer shutdown completed")
471+
except Exception as e:
472+
logger.error(f"Error shutting down SimpleTCPServer: {str(e)}")
500473

474+
# Clear servers list
475+
self.servers = []
476+
477+
# Shutdown lifespan
501478
if self.lifespan is not None:
502479
try:
503480
await self.lifespan.shutdown()
504481
logger.debug("Lifespan shutdown completed")
505482
except Exception as e:
506483
logger.error(f"Error during lifespan shutdown: {str(e)}")
507-
logger.debug(f"Lifespan shutdown error details: {traceback.format_exc()}")
508484

509-
self.servers = []
510485
self.lifespan = None
511486

512487
logger.debug("Server shutdown process completed")
513-
514-
async def serve(self, sockets=None):
515-
self.config.setup_event_loop()
516-
517-
self.lifespan = None
518-
519-
self._initialize_lifespan()
520-
521-
try:
522-
await self.startup(sockets=sockets)
523-
524-
if hasattr(self, 'on_startup_callback') and callable(self.on_startup_callback):
525-
self.on_startup_callback()
526-
527-
await self.main_loop()
528-
await self.shutdown()
529-
except Exception as e:
530-
logger.error(f"Error during server operation: {str(e)}")
531-
logger.debug(f"Server error details: {traceback.format_exc()}")
532-
raise
533-
534-
def _initialize_lifespan(self):
535-
try:
536-
from uvicorn.lifespan.on import LifespanOn
537-
538-
try:
539-
self.lifespan = LifespanOn(config=self.config)
540-
logger.info("Using LifespanOn from uvicorn.lifespan.on with config parameter")
541-
return
542-
except TypeError:
543-
pass
544-
545-
try:
546-
self.lifespan = LifespanOn(self.config.app)
547-
logger.info("Using LifespanOn from uvicorn.lifespan.on with app parameter")
548-
return
549-
except TypeError:
550-
pass
551-
552-
try:
553-
lifespan_on = self.config.lifespan_on if hasattr(self.config, "lifespan_on") else "auto"
554-
self.lifespan = LifespanOn(self.config.app, lifespan_on)
555-
logger.info("Using LifespanOn from uvicorn.lifespan.on with app and lifespan_on parameters")
556-
return
557-
except TypeError:
558-
pass
559-
560-
except (ImportError, AttributeError):
561-
logger.debug("Could not import LifespanOn from uvicorn.lifespan.on")
562-
563-
try:
564-
from uvicorn.lifespan.lifespan import Lifespan
565-
566-
try:
567-
self.lifespan = Lifespan(self.config.app)
568-
logger.info("Using Lifespan from uvicorn.lifespan.lifespan with app parameter")
569-
return
570-
except TypeError:
571-
pass
572-
573-
try:
574-
self.lifespan = Lifespan(self.config.app, "auto")
575-
logger.info("Using Lifespan from uvicorn.lifespan.lifespan with app and lifespan_on parameters")
576-
return
577-
except TypeError:
578-
pass
579-
580-
try:
581-
self.lifespan = Lifespan(self.config.app, "auto", logger)
582-
logger.info("Using Lifespan from uvicorn.lifespan.lifespan with app, lifespan_on, and logger parameters")
583-
return
584-
except TypeError:
585-
pass
586-
587-
except (ImportError, AttributeError):
588-
logger.debug("Could not import Lifespan from uvicorn.lifespan.lifespan")
589-
590-
try:
591-
from uvicorn.lifespan import Lifespan
592-
593-
try:
594-
self.lifespan = Lifespan(self.config.app)
595-
logger.info("Using Lifespan from uvicorn.lifespan with app parameter")
596-
return
597-
except TypeError:
598-
pass
599-
600-
try:
601-
self.lifespan = Lifespan(self.config.app, "auto")
602-
logger.info("Using Lifespan from uvicorn.lifespan with app and lifespan_on parameters")
603-
return
604-
except TypeError:
605-
pass
606-
607-
try:
608-
self.lifespan = Lifespan(self.config.app, "auto", logger)
609-
logger.info("Using Lifespan from uvicorn.lifespan with app, lifespan_on, and logger parameters")
610-
return
611-
except TypeError:
612-
pass
613-
614-
except (ImportError, AttributeError):
615-
logger.debug("Could not import Lifespan from uvicorn.lifespan")
616-
617-
try:
618-
from uvicorn.lifespan.state import LifespanState
619-
620-
try:
621-
self.lifespan = LifespanState(self.config.app)
622-
logger.info("Using LifespanState from uvicorn.lifespan.state with app parameter")
623-
return
624-
except TypeError:
625-
pass
626-
627-
try:
628-
self.lifespan = LifespanState(self.config.app, logger=logger)
629-
logger.info("Using LifespanState from uvicorn.lifespan.state with app and logger parameters")
630-
return
631-
except TypeError:
632-
pass
633-
634-
except (ImportError, AttributeError):
635-
logger.debug("Could not import LifespanState from uvicorn.lifespan.state")
636-
637-
self.lifespan = NoopLifespan(self.config.app)
638-
logger.warning("Using NoopLifespan - server may not handle startup/shutdown events properly")
639-
640488

641489
def start_server(use_ngrok: bool = None, port: int = None, ngrok_auth_token: Optional[str] = None):
642490
try:
@@ -1017,7 +865,7 @@ async def on_startup_async():
1017865
logger.warning("Falling back to direct SimpleTCPServer implementation")
1018866

1019867

1020-
direct_server = SimpleTCPServer(config)
868+
direct_server = SimpleTCPServer(config=self.config)
1021869

1022870

1023871
asyncio.run(direct_server.serve())
@@ -1062,7 +910,7 @@ async def on_startup_async():
1062910
logger.warning("Falling back to direct SimpleTCPServer implementation")
1063911

1064912

1065-
direct_server = SimpleTCPServer(config)
913+
direct_server = SimpleTCPServer(config=self.config)
1066914

1067915

1068916
loop.run_until_complete(direct_server.serve())
@@ -1161,7 +1009,7 @@ async def on_startup_async():
11611009
logger.warning("Falling back to direct SimpleTCPServer implementation")
11621010

11631011

1164-
direct_server = SimpleTCPServer(config)
1012+
direct_server = SimpleTCPServer(config=self.config)
11651013

11661014

11671015
asyncio.run(direct_server.serve())
@@ -1206,7 +1054,7 @@ async def on_startup_async():
12061054
logger.warning("Falling back to direct SimpleTCPServer implementation")
12071055

12081056

1209-
direct_server = SimpleTCPServer(config)
1057+
direct_server = SimpleTCPServer(config=self.config)
12101058

12111059

12121060
loop.run_until_complete(direct_server.serve())

‎setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
setup(
77
name="locallab",
8-
version="0.4.38",
8+
version="0.4.39",
99
packages=find_packages(include=["locallab", "locallab.*"]),
1010
install_requires=[
1111
"fastapi>=0.95.0,<1.0.0",

0 commit comments

Comments
 (0)