Add POET Server

poet-server/Dockerfile

@@ -0,0 +1,20 @@
+# TODO: currently the build context uses the entire POET repo since it installs
+# the poet-ai package locally; it should be isolated in the future
+
+FROM python:3.10-slim
+
+# Install server dependencies
+WORKDIR /app/poet-server
+ADD poet-server/requirements.txt requirements.txt
+RUN pip3 install -r requirements.txt
+
+# Install the poet-ai package
+WORKDIR /app
+ADD setup.py setup.py
+ADD poet poet
+RUN pip3 install -e .
+
+WORKDIR /app/poet-server
+ADD poet-server/server.py server.py
+
+CMD ["python", "server.py"]

poet-server/README.md

@@ -0,0 +1,76 @@
+# POET Server Setup
+
+In this guide, we will show you how to set up the POET ILP server for use locally or as a hosted service. We use this to host our POET server that powers requests from the [POET Demo Colab](https://colab.research.google.com/drive/1iup_edJd9zB1tfVBHXLmkWOT5yoSmXzz?usp=sharing) notebook. You do not need to set-up a POET-server to use POET, but feel free to use it to set-up your hosted service. 
+
+POET's Integer Linear Program (ILP) formulation is compatible with a variety of solvers, including Gurobi and COIN-OR CBC. In this guide, we will demonstrate how to set it up using both of these solvers.
+
+## Setting Up Gurobi (Optional)
+
+The ILP solver defaults to using the COIN-OR CBC solver when Gurobi isn't available. However, since Gurobi is much faster, it is recommended to install it where possible.
+
+### Acquiring a Free Academic Gurobi Web License
+
+1. Create a free Gurobi account [here](https://pages.gurobi.com/registration). Make sure to specify the Academic user option.
+2. Complete the rest of the Gurobi account creation process, which will include creating a password and verifying your email address.
+3. Login to the Gurobi [Web License Manager](https://license.gurobi.com/) using your new account.
+4. Create and download a new Web License file. This will be a `gurobi.lic` file that you will need in later steps, so keep note of where you save it.
+
+## Option 1: Running the Server Locally
+
+1. If using Gurobi, move the `gurobi.lic` file you downloaded in the previous step to your home directory (i.e. to `~/gurobi.lic`).
+2. Clone this repository by running `git clone https://github.com/ShishirPatil/poet`.
+3. Run `pip3 install -e .` in this repository's root directory to install the `poet-ai` package.
+4. Run `cd poet-server` to navigate to the `poet-server` directory.
+5. Run `pip3 install -r requirements.txt` to install the ILP server dependencies.
+6. Finally, run `python3 server.py` to start the server.
+    - You can optionally run `DEV=1 python3 server.py` to enable reload mode, which will automatically restart the server when you make changes to the code.
+7. You can now make requests to the server at `http://localhost/solve`.
+
+## Option 2: Running the Server Locally within a Docker Container
+
+We include a Docker image that can be used to run the server. 
+
+Prebuilt Docker images are available at `public.ecr.aws/i5z6k9k2/poet-server`
+
+You can pull an image and start the server using:
+
+```bash
+docker pull public.ecr.aws/i5z6k9k2/poet-server:latest
+docker run -p 80:80 -v ~/gurobi.lic:/opt/gurobi/gurobi.lic public.ecr.aws/i5z6k9k2/poet-server
+```
+
+Or, you can build the docker container yourself following the steps below. 
+
+
+1. Ensure you have [Docker Compose](https://docs.docker.com/compose/install/) installed.
+2. Clone this repository by running `git clone https://github.com/ShishirPatil/poet`.
+3. If using Gurobi, move the `gurobi.lic` file you downloaded in the previous step to the `poet-server` directory of this repository (i.e. to `poet-server/gurobi.lic`).
+4. Run `cd poet-server` to navigate to the `poet-server` directory.
+5. Run `docker compose up --build` to build and start the Docker container.
+6. You can now make GET requests to the server at `http://localhost/solve` as shown below.
+
+## Option 3: Hosting POET server on an AWS EC2 Instance
+
+Ensure that you have moved the `gurobi.lic` file (if you want to use the Gurobi optimizer) you downloaded earlier to the EC2 instance. Ensure that Port 80 is open for ingress traffic. 
+
+
+## Making Requests 
+
+To issue requests to the POET server, you can use the following Python code. Here, we use the demo POET-server hosted at IP `54.189.43.62`:
+
+```python
+import requests
+
+response = requests.get("http://54.189.43.62/solve", {
+    "model": "linear",
+    "platform": "m0",
+    "ram_budget": 90000000,
+    "runtime_budget": 1.253,
+    "solver": "gurobi",
+})
+
+print(response.json())
+```
+
+
+

poet-server/docker-compose.yml

@@ -0,0 +1,14 @@
+version: "3.3"
+
+services:
+    server:
+        image: public.ecr.aws/i5z6k9k2/poet-server:latest
+        shm_size: 2.5gb
+        tty: true
+        ports:
+            - 80:80
+        build:
+            context: ../
+            dockerfile: poet-server/Dockerfile
+        volumes:
+            - ${PWD}/gurobi.lic:/opt/gurobi/gurobi.lic

poet-server/requirements.txt

@@ -0,0 +1,2 @@
+fastapi
+uvicorn[standard]

poet-server/server.py

@@ -0,0 +1,61 @@
+import os
+from typing import Literal, Optional
+
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from loguru import logger
+
+from poet import solve
+
+SOLVE_THREADS = min(4, os.cpu_count())
+
+app = FastAPI()
+
+
+@app.get("/solve")
+def solve_handler(
+    model: Literal[
+        "linear",
+        "vgg16",
+        "vgg16_cifar",
+        "resnet18",
+        "resnet50",
+        "resnet18_cifar",
+        "bert",
+        "transformer",
+    ],
+    platform: Literal["m0", "a72", "a72nocache", "m4", "jetsontx2"],
+    ram_budget: float,
+    runtime_budget: float,
+    paging: int = 1,
+    remat: int = 1,
+    mem_power_scale=1.0,
+    batch_size=1,
+    use_actual_gurobi: Optional[bool] = None,
+    solver: Optional[Literal["gurobi", "cbc"]] = None,
+    time_limit_s: float = 1e100,
+    solve_threads: int = SOLVE_THREADS,  # different default than a direct solve
+):
+    try:
+        return solve(
+            model=model,
+            platform=platform,
+            ram_budget=ram_budget,
+            runtime_budget=runtime_budget,
+            paging=paging,
+            remat=remat,
+            mem_power_scale=mem_power_scale,
+            batch_size=batch_size,
+            use_actual_gurobi=use_actual_gurobi,
+            solver=solver,
+            time_limit_s=time_limit_s,
+            solve_threads=solve_threads,
+        )
+    except Exception as e:
+        logger.exception(e)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    logger.info("Initializing an instance of the POET server.")
+    uvicorn.run("server:app", host="0.0.0.0", port=80, reload=os.environ.get("DEV"))

poet/util.py

@@ -1,6 +1,6 @@
 import pickle
 from pathlib import Path
-from typing import List, Optional
+from typing import List
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -11,12 +11,11 @@
 from poet.architectures.resnet import resnet18, resnet18_cifar, resnet50
 from poet.architectures.vgg import vgg16
 from poet.chipsets import M4F, MKR1000, JetsonTX2, RPi, RPiNoCache
-from poet.power_computation import DNNLayer, GradientLayer, get_net_costs
 from poet.poet_solver import POETSolution
+from poet.power_computation import DNNLayer, GradientLayer, get_net_costs
 from poet.utils.checkmate.core.dfgraph import DFGraph
 from poet.utils.checkmate.core.graph_builder import GraphBuilder
 from poet.utils.checkmate.core.utils.definitions import PathLike
-from poet.utils.checkmate.core.utils.scheduler import schedule_from_rs
 from poet.utils.checkmate.plot.graph_plotting import plot_dfgraph
 
 
@@ -115,12 +114,14 @@ def plot_network(
 def print_result(result: dict):
     solution: POETSolution = result["solution"]
     if solution.feasible:
-        solution_msg = "successfully found an optimal solution" if solution.finished else "found a feasible solution"
+        solution_msg = "successfully found an optimal solution" if solution.optimal else "found a feasible solution"
         print(
-            f"POET {solution_msg} with a memory budget of {result['ram_budget_bytes']} bytes that consumes {result['total_power_cost_cpu']:.5f} J of CPU power and {result['total_power_cost_page']:.5f} J of memory paging power"
+            f"POET {solution_msg} with a memory budget of {result['ram_budget_bytes']} bytes that consumes {result['total_power_cost_cpu']} J of CPU power and {result['total_power_cost_page']} J of memory paging power"
         )
-        if not solution.finished:
-            print("This solution is not guaranteed to be optimal - you can try increasing the time limit to find an optimal solution")
+        if not solution.optimal:
+            print(
+                "This solution is not guaranteed to be optimal - you can try increasing the solve time [time_limit_s] to find an optimal solution"
+            )
 
         plt.matshow(solution.R)
         plt.title("R")
@@ -133,54 +134,7 @@ def print_result(result: dict):
         plt.matshow(solution.SSd)
         plt.title("SSd")
         plt.show()
-
-        plot_schedule(solution.R, solution.SRam, show=True, plot_mem_usage=True)
-    elif solution.finished:
-        print("POET finished solving and determined that no feasible solution exists")
-    else:
-        print("POET failed to find a feasible solution within the provided time limit")
-
-
-def plot_schedule(R, S, plot_mem_usage=False, mem_grid=None, U=None, save_file: Optional[PathLike] = None, show=False):
-    x, y = get_chipset_and_net("m0", "linear", 1, 1)
-    g, *_ = make_dfgraph_costs(y, x)
-    _, scheduler_aux_data = schedule_from_rs(g, np.array(R), np.array(S))
-    mem_grid = scheduler_aux_data.mem_grid
-
-    if plot_mem_usage:
-        # assert mem_grid is not None
-        fig, axs = plt.subplots(1, 4)
-        vmax = mem_grid
-        vmax = vmax if U is None else max(vmax, np.max(U))
-
-        # Plot slow verifier memory usage
-        axs[2].invert_yaxis()
-        axs[2].pcolormesh(mem_grid, cmap="Greys", vmin=0, vmax=vmax[0][0])
-        axs[2].set_title("Memory usage (verifier)")
-
-        # Plot solver memory usage variables
-        axs[3].invert_yaxis()
-        axs[3].set_title("Memory usage (solved)")
-        if U is not None:
-            axs[3].pcolormesh(U, cmap="Greys", vmin=0, vmax=vmax)
-
-        fig.set_size_inches(28, 6)
     else:
-        fig, axs = plt.subplots(1, 2)
-        fig.set_size_inches(18, 6)
-
-    axs[0].invert_yaxis()
-    axs[0].pcolormesh(R, cmap="Greys", vmin=0, vmax=1)
-    axs[0].set_title("R")
-
-    axs[1].invert_yaxis()
-    axs[1].pcolormesh(S, cmap="Greys", vmin=0, vmax=1)
-    axs[1].set_title("S")
-
-    if show:
-        plt.show()
-    if save_file:
-        path = Path(save_file)
-        path.parents[0].mkdir(parents=True, exist_ok=True)
-        fig.savefig(path)
-        plt.close(fig)
+        print(
+            "POET failed to find a feasible solution within the provided time limit. \n Either a) increase the memory and training time budgets, and/or b) increase the solve time [total_power_cost_page]"
+        )