centml 0.1.0__py3-none-any.whl

centml/__init__.py

@@ -0,0 +1 @@
+from centml.compiler.main import compile

centml/cli/__init__.py

@@ -0,0 +1 @@
+from centml.cli.main import cli, ccluster

centml/cli/cluster.py

@@ -0,0 +1,192 @@
+import sys
+from functools import wraps
+from typing import Dict
+import click
+from tabulate import tabulate
+from centml.sdk import DeploymentType, DeploymentStatus, HealthStatus, ApiException, HardwareInstanceResponse
+from centml.sdk.api import get_centml_client
+
+
+depl_name_to_type_map = {
+    "inference": DeploymentType.INFERENCE_V2,
+    "compute": DeploymentType.COMPUTE_V2,
+    "cserve": DeploymentType.CSERVE,
+}
+depl_type_to_name_map = {v: k for k, v in depl_name_to_type_map.items()}
+
+
+def handle_exception(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except ApiException as e:
+            click.echo(f"Error: {e.reason}")
+            return None
+
+    return wrapper
+
+
+def _get_hw_to_id_map(cclient, cluster_id):
+    response = cclient.get_hardware_instances(cluster_id)
+
+    # Initialize hashmap for hardware to id or vice versa mapping
+    hw_to_id_map: Dict[str, int] = {}
+    id_to_hw_map: Dict[int, HardwareInstanceResponse] = {}
+
+    for hw in response:
+        hw_to_id_map[hw.name] = hw.id
+        id_to_hw_map[hw.id] = hw
+    return hw_to_id_map, id_to_hw_map
+
+
+def _format_ssh_key(ssh_key):
+    if not ssh_key:
+        return "No SSH Key Found"
+    return ssh_key[:10] + '...'
+
+
+def _get_ready_status(cclient, deployment):
+    api_status = deployment.status
+    service_status = (
+        cclient.get_status(deployment.id).service_status if deployment.status == DeploymentStatus.ACTIVE else None
+    )
+
+    status_styles = {
+        (DeploymentStatus.PAUSED, None): ("paused", "yellow", "black"),
+        (DeploymentStatus.DELETED, None): ("deleted", "white", "black"),
+        (DeploymentStatus.ACTIVE, HealthStatus.HEALTHY): ("ready", "green", "black"),
+        (DeploymentStatus.ACTIVE, HealthStatus.PROGRESSING): ("starting", "black", "white"),
+        (DeploymentStatus.ACTIVE, HealthStatus.DEGRADED): ("starting", "black", "white"),
+        (DeploymentStatus.ACTIVE, HealthStatus.MISSING): ("not found", "cyan"),
+    }
+
+    style = status_styles.get((api_status, service_status), ("unknown", "black", "white"))
+    # Handle foreground and background colors
+    return click.style(style[0], fg=style[1], bg=style[2])
+
+
+@click.command(help="List all deployments")
+@click.argument("type", type=click.Choice(list(depl_name_to_type_map.keys())), required=False, default=None)
+def ls(type):
+    with get_centml_client() as cclient:
+        depl_type = depl_name_to_type_map[type] if type in depl_name_to_type_map else None
+        deployments = cclient.get(depl_type)
+        rows = [
+            [d.id, d.name, depl_type_to_name_map[d.type], d.status.value, d.created_at.strftime("%Y-%m-%d %H:%M:%S")]
+            for d in deployments
+        ]
+
+        click.echo(
+            tabulate(
+                rows,
+                headers=["ID", "Name", "Type", "Status", "Created at"],
+                tablefmt="rounded_outline",
+                disable_numparse=True,
+            )
+        )
+
+
+@click.command(help="Get deployment details")
+@click.argument("type", type=click.Choice(list(depl_name_to_type_map.keys())))
+@click.argument("id", type=int)
+@handle_exception
+def get(type, id):
+    with get_centml_client() as cclient:
+        depl_type = depl_name_to_type_map[type]
+
+        if depl_type == DeploymentType.INFERENCE_V2:
+            deployment = cclient.get_inference(id)
+        elif depl_type == DeploymentType.COMPUTE_V2:
+            deployment = cclient.get_compute(id)
+        elif depl_type == DeploymentType.CSERVE:
+            deployment = cclient.get_cserve(id)
+        else:
+            sys.exit("Please enter correct deployment type")
+
+        ready_status = _get_ready_status(cclient, deployment)
+        _, id_to_hw_map = _get_hw_to_id_map(cclient, deployment.cluster_id)
+        hw = id_to_hw_map[deployment.hardware_instance_id]
+
+        click.echo(
+            tabulate(
+                [
+                    ("Name", deployment.name),
+                    ("Status", ready_status),
+                    ("Endpoint", deployment.endpoint_url),
+                    ("Created at", deployment.created_at.strftime("%Y-%m-%d %H:%M:%S")),
+                    ("Hardware", f"{hw.name} ({hw.num_gpu}x {hw.gpu_type})"),
+                    ("Cost", f"{hw.cost_per_hr/100} credits/hr"),
+                ],
+                tablefmt="rounded_outline",
+                disable_numparse=True,
+            )
+        )
+
+        click.echo("Additional deployment configurations:")
+        if depl_type == DeploymentType.INFERENCE_V2:
+            click.echo(
+                tabulate(
+                    [
+                        ("Image", deployment.image_url),
+                        ("Container port", deployment.container_port),
+                        ("Healthcheck", deployment.healthcheck or "/"),
+                        ("Replicas", {"min": deployment.min_scale, "max": deployment.max_scale}),
+                        ("Environment variables", deployment.env_vars or "None"),
+                        ("Max concurrency", deployment.concurrency or "None"),
+                    ],
+                    tablefmt="rounded_outline",
+                    disable_numparse=True,
+                )
+            )
+        elif depl_type == DeploymentType.COMPUTE_V2:
+            click.echo(
+                tabulate(
+                    [("Username", "centml"), ("SSH key", _format_ssh_key(deployment.ssh_public_key))],
+                    tablefmt="rounded_outline",
+                    disable_numparse=True,
+                )
+            )
+        elif depl_type == DeploymentType.CSERVE:
+            click.echo(
+                tabulate(
+                    [
+                        ("Hugging face model", deployment.model),
+                        (
+                            "Parallelism",
+                            {"tensor": deployment.tensor_parallel_size, "pipeline": deployment.pipeline_parallel_size},
+                        ),
+                        ("Replicas", {"min": deployment.min_scale, "max": deployment.max_scale}),
+                        ("Max concurrency", deployment.concurrency or "None"),
+                    ],
+                    tablefmt="rounded_outline",
+                    disable_numparse=True,
+                )
+            )
+
+
+@click.command(help="Delete a deployment")
+@click.argument("id", type=int)
+@handle_exception
+def delete(id):
+    with get_centml_client() as cclient:
+        cclient.delete(id)
+        click.echo("Deployment has been deleted")
+
+
+@click.command(help="Pause a deployment")
+@click.argument("id", type=int)
+@handle_exception
+def pause(id):
+    with get_centml_client() as cclient:
+        cclient.pause(id)
+        click.echo("Deployment has been paused")
+
+
+@click.command(help="Resume a deployment")
+@click.argument("id", type=int)
+@handle_exception
+def resume(id):
+    with get_centml_client() as cclient:
+        cclient.resume(id)
+        click.echo("Deployment has been resumed")

centml/cli/login.py

@@ -0,0 +1,30 @@
+import click
+
+from centml.sdk import auth
+from centml.sdk.config import settings
+
+
+@click.command(help="Login to CentML")
+@click.argument("token_file", required=False)
+def login(token_file):
+    if token_file:
+        auth.store_centml_cred(token_file)
+
+    if auth.load_centml_cred():
+        click.echo(f"Authenticating with credentials from {settings.CENTML_CRED_FILE_PATH}\n")
+        click.echo("Login successful")
+    else:
+        click.echo("Login with CentML authentication token")
+        click.echo("Usage: centml login TOKEN_FILE\n")
+        choice = click.confirm("Do you want to download the token?")
+
+        if choice:
+            click.launch(f"{settings.CENTML_WEB_URL}?isCliAuthenticated=true")
+        else:
+            click.echo("Login unsuccessful")
+
+
+@click.command(help="Logout from CentML")
+def logout():
+    auth.remove_centml_cred()
+    click.echo("Logout successful")

centml/cli/main.py

@@ -0,0 +1,35 @@
+import click
+
+from centml.cli.login import login, logout
+from centml.cli.cluster import ls, get, delete, pause, resume
+
+
+@click.group()
+def cli():
+    pass
+
+
+cli.add_command(login)
+cli.add_command(logout)
+
+
+@cli.command(help="Start remote compilation server")
+def server():
+    from centml.compiler.server import run
+
+    run()
+
+
+@click.group(help="CentML cluster CLI tool")
+def ccluster():
+    pass
+
+
+ccluster.add_command(ls)
+ccluster.add_command(get)
+ccluster.add_command(delete)
+ccluster.add_command(pause)
+ccluster.add_command(resume)
+
+
+cli.add_command(ccluster, name="cluster")

centml/compiler/__init__.py

@@ -0,0 +1,3 @@
+from centml.compiler.main import compile
+
+all = ["compile"]

centml/compiler/backend.py

@@ -0,0 +1,194 @@
+import os
+import gc
+import time
+import hashlib
+import logging
+import threading as th
+from http import HTTPStatus
+from weakref import ReferenceType, ref
+from tempfile import TemporaryDirectory
+from typing import List, Callable, Optional
+import requests
+import torch
+from torch.fx import GraphModule
+from centml.compiler.config import settings, CompilationStatus
+from centml.compiler.utils import get_backend_compiled_forward_path
+
+
+class Runner:
+    def __init__(self, module: GraphModule, inputs: List[torch.Tensor]):
+        if not module:
+            raise Exception("No module provided")
+
+        self._module: ReferenceType[GraphModule] = ref(module)
+        self._inputs: List[torch.Tensor] = inputs
+        self.compiled_forward_function: Optional[Callable[[torch.Tensor], tuple]] = None
+        self.lock = th.Lock()
+        self.child_thread = th.Thread(target=self.remote_compilation_starter)
+
+        self.serialized_model_dir: Optional[TemporaryDirectory] = None
+        self.serialized_model_path: Optional[str] = None
+        self.serialized_input_path: Optional[str] = None
+
+        try:
+            self.child_thread.start()
+        except Exception as e:
+            logging.getLogger(__name__).exception(f"Failed to start compilation thread\n{e}")
+
+    @property
+    def module(self) -> Optional[GraphModule]:
+        return self._module()
+
+    @module.deleter
+    def module(self):
+        self._module().graph.owning_module = None
+        self._module = None
+
+    @property
+    def inputs(self) -> List[torch.Tensor]:
+        return self._inputs
+
+    @inputs.deleter
+    def inputs(self):
+        self._inputs = None
+
+    def _serialize_model_and_inputs(self):
+        self.serialized_model_dir = TemporaryDirectory()  # pylint: disable=consider-using-with
+        self.serialized_model_path = os.path.join(self.serialized_model_dir.name, settings.CENTML_SERIALIZED_MODEL_FILE)
+        self.serialized_input_path = os.path.join(self.serialized_model_dir.name, settings.CENTML_SERIALIZED_INPUT_FILE)
+
+        # torch.save saves a zip file full of pickled files with the model's states.
+        try:
+            torch.save(self.module, self.serialized_model_path, pickle_protocol=settings.CENTML_PICKLE_PROTOCOL)
+            torch.save(self.inputs, self.serialized_input_path, pickle_protocol=settings.CENTML_PICKLE_PROTOCOL)
+        except Exception as e:
+            raise Exception(f"Failed to save module or inputs with torch.save: {e}") from e
+
+    def _get_model_id(self) -> str:
+        if not self.serialized_model_path or not os.path.isfile(self.serialized_model_path):
+            raise Exception(f"Model not saved at path {self.serialized_model_path}")
+
+        sha_hash = hashlib.sha256()
+        with open(self.serialized_model_path, "rb") as serialized_model_file:
+            # Read in chunks to not load too much into memory
+            for block in iter(lambda: serialized_model_file.read(settings.CENTML_HASH_CHUNK_SIZE), b""):
+                sha_hash.update(block)
+
+        model_id = sha_hash.hexdigest()
+        logging.info(f"Model has id {model_id}")
+        return model_id
+
+    def _download_model(self, model_id: str):
+        download_response = requests.get(
+            url=f"{settings.CENTML_SERVER_URL}/download/{model_id}", timeout=settings.CENTML_COMPILER_TIMEOUT
+        )
+        if download_response.status_code != HTTPStatus.OK:
+            raise Exception(
+                f"Download: request failed, exception from server:\n{download_response.json().get('detail')}"
+            )
+        if download_response.content == b"":
+            raise Exception("Download: empty response from server")
+        download_path = get_backend_compiled_forward_path(model_id)
+        with open(download_path, "wb") as f:
+            f.write(download_response.content)
+        return torch.load(download_path)
+
+    def _compile_model(self, model_id: str):
+        # The model should have been saved using torch.save when we found the model_id
+        if not self.serialized_model_path or not self.serialized_input_path:
+            raise Exception("Model or inputs not serialized")
+        if not os.path.isfile(self.serialized_model_path):
+            raise Exception(f"Model not saved at path {self.serialized_model_path}")
+        if not os.path.isfile(self.serialized_input_path):
+            raise Exception(f"Inputs not saved at path {self.serialized_input_path}")
+
+        with open(self.serialized_model_path, 'rb') as model_file, open(self.serialized_input_path, 'rb') as input_file:
+            compile_response = requests.post(
+                url=f"{settings.CENTML_SERVER_URL}/submit/{model_id}",
+                files={"model": model_file, "inputs": input_file},
+                timeout=settings.CENTML_COMPILER_TIMEOUT,
+            )
+        if compile_response.status_code != HTTPStatus.OK:
+            raise Exception(
+                f"Compile model: request failed, exception from server:\n{compile_response.json().get('detail')}\n"
+            )
+
+    def _wait_for_status(self, model_id: str) -> bool:
+        tries = 0
+        while True:
+            # get server compilation status
+            status = None
+            try:
+                status_response = requests.get(
+                    f"{settings.CENTML_SERVER_URL}/status/{model_id}", timeout=settings.CENTML_COMPILER_TIMEOUT
+                )
+                if status_response.status_code != HTTPStatus.OK:
+                    raise Exception(
+                        f"Status check: request failed, exception from server:\n{status_response.json().get('detail')}"
+                    )
+                status = status_response.json().get("status")
+            except Exception as e:
+                logging.getLogger(__name__).exception(f"Status check failed:\n{e}")
+
+            if status == CompilationStatus.DONE.value:
+                return True
+            elif status == CompilationStatus.COMPILING.value:
+                pass
+            elif status == CompilationStatus.NOT_FOUND.value:
+                logging.info("Submitting model to server for compilation.")
+                try:
+                    self._compile_model(model_id)
+                except Exception as e:
+                    logging.getLogger(__name__).exception(f"Submitting compilation failed:\n{e}")
+                tries += 1
+            else:
+                tries += 1
+
+            if tries > settings.CENTML_COMPILER_MAX_RETRIES:
+                raise Exception("Waiting for status: compilation failed too many times.\n")
+
+            time.sleep(settings.CENTML_COMPILER_SLEEP_TIME)
+
+    def remote_compilation_starter(self):
+        try:
+            self.remote_compilation()
+        except Exception as e:
+            logging.getLogger(__name__).exception(f"Compilation thread failed:\n{e}")
+
+    def remote_compilation(self):
+        self._serialize_model_and_inputs()
+
+        model_id = self._get_model_id()
+
+        # check if compiled forward is saved locally
+        compiled_forward_path = get_backend_compiled_forward_path(model_id)
+        if os.path.isfile(compiled_forward_path):
+            logging.info("Compiled model found in local cache. Not submitting to server.")
+            compiled_forward = torch.load(compiled_forward_path)
+        else:
+            self._wait_for_status(model_id)
+            compiled_forward = self._download_model(model_id)
+
+        self.compiled_forward_function = compiled_forward
+
+        logging.info("Compilation successful.")
+
+        # Let garbage collector free the memory used by the uncompiled model
+        with self.lock:
+            del self.inputs
+            if self.module:
+                del self.module
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    def __call__(self, *args, **kwargs):
+        # If model is currently compiling, return the uncompiled forward function
+        with self.lock:
+            if not self.compiled_forward_function:
+                return self.module.forward(*args, **kwargs)
+
+        return self.compiled_forward_function(*args)
+
+
+def centml_dynamo_backend(gm: GraphModule, example_inputs: List[torch.Tensor]):
+    return Runner(gm, example_inputs)

centml/compiler/config.py

@@ -0,0 +1,45 @@
+import os
+from enum import Enum
+from pydantic_settings import BaseSettings
+
+
+class CompilationStatus(Enum):
+    NOT_FOUND = "NOT_FOUND"
+    COMPILING = "COMPILING"
+    DONE = "DONE"
+
+
+class OperationMode(Enum):
+    PREDICTION = "PREDICTION"
+    REMOTE_COMPILATION = "REMOTE_COMPILATION"
+
+
+class Config(BaseSettings):
+    CENTML_COMPILER_TIMEOUT: int = 10
+    CENTML_COMPILER_MAX_RETRIES: int = 3
+    CENTML_COMPILER_SLEEP_TIME: int = 15
+
+    CENTML_BASE_CACHE_DIR: str = os.path.expanduser("~/.cache/centml")
+    CENTML_BACKEND_BASE_PATH: str = os.path.join(CENTML_BASE_CACHE_DIR, "backend")
+    CENTML_SERVER_BASE_PATH: str = os.path.join(CENTML_BASE_CACHE_DIR, "server")
+
+    CENTML_SERVER_URL: str = "http://0.0.0.0:8090"
+
+    # Use a constant path since torch.save uses the given file name in it's zipfile.
+    # Using a different filename would result in a different hash.
+    CENTML_SERIALIZED_MODEL_FILE: str = "serialized_model.zip"
+    CENTML_SERIALIZED_INPUT_FILE: str = "serialized_input.zip"
+    CENTML_PICKLE_PROTOCOL: int = 4
+
+    CENTML_HASH_CHUNK_SIZE: int = 4096
+
+    # If the server response is smaller than this, don't gzip it
+    CENTML_MINIMUM_GZIP_SIZE: int = 1000
+
+    CENTML_MODE: OperationMode = OperationMode.REMOTE_COMPILATION
+    CENTML_PREDICTION_DATA_FILE: str = 'tests/sample_data.csv'
+    CENTML_PREDICTION_GPUS: str = "A10G,A100SXM440GB"
+    CENTML_PROMETHEUS_PORT: int = 8000
+
+
+settings = Config()

centml/compiler/main.py

@@ -0,0 +1,57 @@
+import builtins
+from typing import Callable, Dict, Optional, Union
+
+from centml.compiler.config import OperationMode, settings
+
+
+def compile(
+    model: Optional[Callable] = None,
+    *,
+    fullgraph: builtins.bool = False,
+    dynamic: Optional[builtins.bool] = None,
+    mode: Union[str, None] = None,
+    options: Optional[Dict[str, Union[str, builtins.int, builtins.bool]]] = None,
+    disable: builtins.bool = False,
+) -> Callable:
+    import torch
+
+    if settings.CENTML_MODE == OperationMode.REMOTE_COMPILATION:
+        from centml.compiler.backend import centml_dynamo_backend
+
+        # Return the remote-compiled model
+        compiled_model = torch.compile(
+            model,
+            backend=centml_dynamo_backend,  # Compilation backend
+            fullgraph=fullgraph,
+            dynamic=dynamic,
+            mode=mode,
+            options=options,
+            disable=disable,
+        )
+        return compiled_model
+    elif settings.CENTML_MODE == OperationMode.PREDICTION:
+        from centml.compiler.prediction.backend import centml_prediction_backend, get_gauge
+
+        # Proceed with prediction workflow
+        compiled_model = torch.compile(
+            model,
+            backend=centml_prediction_backend,  # Prediction backend
+            fullgraph=fullgraph,
+            dynamic=dynamic,
+            mode=mode,
+            options=options,
+            disable=disable,
+        )
+
+        def centml_wrapper(*args, **kwargs):
+            out = compiled_model(*args, **kwargs)
+            # Update the prometheus metrics with final values
+            gauge = get_gauge()
+            for gpu in settings.CENTML_PREDICTION_GPUS.split(','):
+                gauge.set_metric_value(gpu)
+
+            return out
+
+        return centml_wrapper
+    else:
+        raise Exception("Invalid operation mode")

centml/compiler/prediction/backend.py

@@ -0,0 +1,28 @@
+from typing import List
+
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+from centml.compiler.config import settings
+from centml.compiler.prediction.kdtree import get_tree_db
+from centml.compiler.prediction.metric import get_gauge
+from centml.compiler.prediction.profiler import Profiler
+
+
+def centml_prediction_backend(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
+    profilers = []
+    tree_db = get_tree_db()
+    for gpu in settings.CENTML_PREDICTION_GPUS.split(','):
+        profilers.append(Profiler(gm, gpu, tree_db))
+
+    def forward(*args):
+        fake_mode = FakeTensorMode(allow_non_fake_inputs=True)
+        fake_args = [fake_mode.from_tensor(arg) if isinstance(arg, torch.Tensor) else arg for arg in args]
+        with fake_mode:
+            for prof in profilers:
+                out, t = prof.propagate(*fake_args)
+                gauge = get_gauge()
+                gauge.increment(prof.gpu, t)
+        return out
+
+    return forward

centml/compiler/prediction/kdtree.py

@@ -0,0 +1,69 @@
+import ast
+import csv
+import logging
+
+from sklearn.neighbors import KDTree  # type: ignore
+
+from centml.compiler.config import settings
+
+_tree_db = None
+
+
+class KDTreeWithValues:
+    def __init__(self, points=None, values=None):
+        self.points = points if points else []
+        self.values = values if values else []
+        if self.points:
+            self.tree = KDTree(self.points)
+        else:
+            self.tree = None
+
+    def add(self, point, value):
+        self.points.append(point)
+        self.values.append(value)
+        self.tree = KDTree(self.points)
+
+    def query(self, point):
+        if self.tree is None:
+            return None, None
+
+        dist, idx = self.tree.query([point], k=1)
+        return dist[0][0], self.values[idx[0][0]]
+
+
+class TreeDB:
+    def __init__(self, data_csv):
+        self.db = {}
+        self._populate_db(data_csv)
+
+    def get(self, key, inp):
+        if key not in self.db:
+            logging.getLogger(__name__).warning(f"Key {key} not found in database")
+            return float('-inf')
+            # TODO: Handle the case of unfound keys better. For now, return -inf to indicate something went wrong.
+            # Ideally, we shouldn't throw away a whole prediction because of one possibly insignificant node.
+
+        _, val = self.db[key].query(inp)
+        return val
+
+    def _add_from_db(self, key, points, values):
+        self.db[key] = KDTreeWithValues(points, values)
+
+    def _populate_db(self, data_csv):
+        with open(data_csv, newline='') as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                try:
+                    key = (row['op'], int(row['dim']), row['inp_dtypes'], row['out_dtypes'], row['gpu'])
+                    points = ast.literal_eval(row['points'])
+                    values = ast.literal_eval(row['values'])
+                    self._add_from_db(key, points, values)
+                except ValueError as e:
+                    logging.getLogger(__name__).exception(f"Error parsing row: {row}\n{e}")
+
+
+def get_tree_db():
+    global _tree_db
+    if _tree_db is None:
+        _tree_db = TreeDB(settings.CENTML_PREDICTION_DATA_FILE)
+    return _tree_db