File size: 1,013 Bytes
cfc79aa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import os
import subprocess
import sys
from pathlib import Path

# Main entry point for the multi-node Hugging Face API system
# This can be used to start either the head node or a worker node based on environment

def main():
    node_type = os.getenv("NODE_TYPE", "head")  # Default to head node
    
    if node_type == "head":
        print("πŸš€ Starting Head Node API Gateway...")
        from head_node.app import app
        import uvicorn
        uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))
    elif node_type in ["nano", "mini", "fast", "large"]:
        print(f"πŸš€ Starting Worker Node for {node_type} model...")
        from worker_nodes.worker_app import app
        import uvicorn
        uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 8000)))
    else:
        print(f"Unknown node type: {node_type}")
        print("Use NODE_TYPE=head for head node, or NODE_TYPE=[nano|mini|fast|large] for workers")
        sys.exit(1)

if __name__ == "__main__":
    main()