Refactoring centml modules (#23)

* Rename centml to ccompute
* Change remote compilation server name to "compile-server"
* Move config_instance and CompilationStatus to config.py
* Move centml compiler backend registration to centml.compiler
* Add string annotations to CompilationStatus enums

Author: anandj91

bin/centml bin/ccomputebin/centml renamed to bin/ccompute

@@ -1,9 +0,0 @@
-import hidet
-import torch._dynamo
-from .compiler import backend
-
-hidet.option.imperative(False)
-
-# Register centml compiler backend to torch dynamo
-if "centml" not in torch._dynamo.list_backends():
-    torch._dynamo.register_backend(compiler_fn=backend.centml_dynamo_backend, name="centml")

centml/__init__.py

@@ -5,11 +5,11 @@
 
 def main():
     parser = argparse.ArgumentParser(description="CentML command line tool")
-    parser.add_argument("mode", choices=["server"], type=str)
+    parser.add_argument("mode", choices=["compile-server"], type=str)
 
     args = parser.parse_args()
 
-    if args.mode == "server":
+    if args.mode == "compile-server":
         server.run()
     else:
         raise Exception("Invalid mode")

centml/cli/main.py

@@ -1,3 +1,8 @@
-from .config import Config
+import torch._dynamo
 
-config_instance = Config()
+from .backend import centml_dynamo_backend
+
+
+# Register centml compiler backend to torch dynamo
+if "centml" not in torch._dynamo.list_backends():
+    torch._dynamo.register_backend(compiler_fn=centml_dynamo_backend, name="centml")

centml/compiler/__init__.py

@@ -15,10 +15,10 @@
 from hidet.graph.frontend.torch.interpreter import Interpreter
 from hidet.graph.frontend.torch.dynamo_backends import get_flow_graph, get_wrapper
 from hidet.runtime.compiled_graph import load_compiled_graph, CompiledGraph
-from centml.compiler import config_instance
-from centml.compiler.server_compilation import CompilationStatus
+from centml.compiler.config import config_instance, CompilationStatus
 
 
+hidet.option.imperative(False)
 base_path = os.path.join(config_instance.CACHE_PATH, "compiler")
 os.makedirs(base_path, exist_ok=True)
 server_url = f"http://{config_instance.SERVER_IP}:{config_instance.SERVER_PORT}"

centml/compiler/backend.py

@@ -1,4 +1,11 @@
 import os
+from enum import Enum
+
+
+class CompilationStatus(Enum):
+    NOT_FOUND = "not_found"
+    COMPILING = "compiling"
+    DONE = "done"
 
 
 class Config:
@@ -9,3 +16,6 @@ class Config:
     CACHE_PATH = os.getenv("CENTML_CACHE_DIR", default=os.path.expanduser("~/.cache/centml"))
     SERVER_IP = os.getenv("CENTML_SERVER_IP", default="0.0.0.0")
     SERVER_PORT = os.getenv("CENTML_SERVER_PORT", default="8080")
+
+
+config_instance = Config()

centml/compiler/config.py

@@ -6,8 +6,8 @@
 from fastapi import FastAPI, UploadFile, HTTPException, BackgroundTasks, Response
 from fastapi.responses import FileResponse
 from fastapi.middleware.gzip import GZipMiddleware
-from centml.compiler.server_compilation import hidet_backend_server, storage_path, CompilationStatus, dir_cleanup
-from centml.compiler import config_instance
+from centml.compiler.server_compilation import hidet_backend_server, storage_path, dir_cleanup
+from centml.compiler.config import config_instance, CompilationStatus
 
 logger = logging.getLogger(__name__)
 

centml/compiler/server.py

@@ -1,21 +1,14 @@
 import os
 import shutil
 import logging
-from enum import Enum
 from typing import List
 from torch import Tensor
 from torch.fx import GraphModule
 from hidet.runtime.compiled_graph import save_compiled_graph
 from hidet.graph.frontend import from_torch
 from hidet.graph.frontend.torch.interpreter import Interpreter
 from hidet.graph.frontend.torch.dynamo_backends import get_flow_graph, get_compiled_graph, preprocess_inputs
-from centml.compiler import config_instance
-
-
-class CompilationStatus(Enum):
-    NOT_FOUND = 1
-    COMPILING = 2
-    DONE = 3
+from centml.compiler.config import config_instance
 
 
 storage_path = os.path.join(config_instance.CACHE_PATH, "server")

centml/compiler/server_compilation.py

@@ -14,7 +14,7 @@
     name='centml',
     version='0.1.0',
     packages=find_packages(),
-    scripts=['bin/centml'],
-	install_requires=REQUIRES
+    scripts=['bin/ccompute'],
+    install_requires=REQUIRES
 )
 

setup.py

@@ -7,8 +7,7 @@
 import hidet
 from hidet.graph.frontend.torch.dynamo_backends import get_flow_graph
 from centml.compiler.backend import Runner
-from centml.compiler.server_compilation import CompilationStatus
-from centml.compiler import config_instance
+from centml.compiler.config import config_instance, CompilationStatus
 
 
 class TestGetModelId(TestCase):