Built a Distributed multi VMs system

this article is solution of this assignment Alchemyst-ai

We are given an AI system which we can run a single system in a single process?

User request -> model -> response

But what problems we can have?

• ▸What if system becomes overloaded?
• ▸What if we wanna scale it?
• ▸What if we wanna run 100x larger model?
• ▸What if different parts of code is written in different languages?

So, instead of having everything in a single system we will split responsibilies.

• ▸For example here, Machine 1 for API and Machine 2 for Model Worker

But one question might be coming to your mind like if we put things on a different machines then how will talk?

By network communication!!
by invoking functions remotely (Remote Procedure Call[RPC])

Now, let's see how the flow would look like.464

Now, think - do you expose internal machines publicly? like do you wants that anyone would come and make any no. of requests to our model?

No right. So what we should do? But before that let's understand what all problems we can have if we happens to do so.

• ▸huge amounts of bills thrown to your door.
• ▸no authentication.
• ▸model may die.
• ▸Solution - Use Gateways: public internet -> gateway only -> private subnet workers

Now, let's explore the two workers first

caller-worker or api-worker(Typescript)

Function:

inference::get_response

One question might be coming to your mind like couldn't HTTP directly call Python worker? yes, obviously. we can do something like:- http request -> python worker -> response

But, what are security? Like what if we wanna add:-

• ▸authentication (only registered users can access)
• ▸logging (every request and response is stored)
• ▸rate limiting (Limit the no. of requests in given timeframe)
• ▸retries
• ▸request validation (are we sending right request?)
• ▸orchestration between many workers

So that's why caller worker and it's flow looks like HTTP request -> caller-worker -> python inference worker -> response

Inference-worker (Python)

Function:

inference::run_inference

What does this do? This is the actual AI brain.

• ▸Flow:- messages -> gemma model -> generated response
• ▸Python-woker: load model, run inference, return answer

Like:- Input

{
  "messages": [
    {
      "role": "user",
      "content": "Explain Redis simply"
    }
  ]
}

Caller-worker HTTP endpoint

Function:

http::run_inference_over_http

But why do we even need this function? Because model workers don't speak “browser”. Users talk HTTP:

POST /v1/chat/completions

Workers talk RPC. So we need translation.

Full request Cycle

User sends:

POST /v1/chat/completions

with:

{
  "messages": [
    {
      "role": "user",
      "content": "What is Redis?"
    }
  ]
}

1. HTTP endpoint receiver request. http::run_inference_over_http. Now it forwards this request.

Loading diagram…

2. Typescript caller worker receives. inference::get_response. It asks inference worker for response. So makes RPC call to inference::run_inference

3. Python worker runs model: gemma model. Returns:-

"Redis is an in-memory data store..."

4. Response bubbles back to user. Final response

{
  "response": "Redis is an in-memory datastore..."
}

Architecture

Loading diagram…

Now before moving forward let's understand some jargons first

• ▸process? running code is process. every process has: memory, CPU usage, ports, lifecycle
• ▸And our architecture is like:- Process A <---network--->Process B. We can say it as distrubed systems at smallest scale.
• ▸Port? Think apartment building. Machine IP = building address. Port = apartment number.
• ▸iii-workers: self contained modular processes in the iii engine; basically like robots(iii workers) working in the factory (iii engine)
• ▸IaC(Infrastructure as Code): using code to manage the infrastructure like building the whole system in aws by code. Terraform ia an open-source IaC

Example:

localhost:3000
localhost:5000

Same machine.

Different services.

Flow Revisited

Loading diagram…

Loading diagram…

Now, let's visualize how the execution flow looks like

Loading diagram…

And the expected logs would be:

Caller worker started
Inference worker started
http server listening

Okay now let's move onto understanding and deploying it into AWS

• ▸First we will create a VPC (Virtual Private Cloud). Think it like a building where many apartments(VMs) would exist.
• ▸then we will create subnets. subnets are like logical divisions of IP network like it defines the range within which a single type will exist like public or private.
• ▸then Internet Gateway and NAT Gateway: Internet Gateway - inbound and outbound internet both; NAT Gateway - outbound internet
• ▸Route Tables: It control the direction of traffic of subnets like where they will go, else they will be clueless
• ▸Security Groups: Firewall Rules(Allows or blocks IPs)
• ▸API VM and Inference VM(Virtual Machine) or 2 EC2 Machines

Let's bit talk about security

What do you think should our model be accessible to everyone? no right? so we will place it in private subnet.

But you me be thinking then how we are gonna access internet for downloading our models and dependencies? as in private subnet there is no public IP right. Yeah you are right. We are gonna use NAT gateways for that purpose we are gonna request it whatever we needs and it will bring that from the internet for us.

Again another question you might be thinking how we are gonna SSH into it in the first place? We can't directly SSH into it but we can do that via the API VM

like in building there are many apartments and accessible to society peoples. similarly in vpc we can jump from one private vm to another this is called bastion jump. like:- internet ---> API VM ---> inference VM

Now, one question must be popping into you mind like what if someone knows our password then they can enter? yeah so for that purpose we can use firewall / security groups. and lock the SSH Port for only our IP

Loading diagram…

Implementation: https://github.com/paramcodes/devops-infra