Add CENTML_ to start of config vars (#74)

* add CENTML to config vars

Author: destefy

centml/compiler/backend.py

@@ -54,13 +54,13 @@ def inputs(self):
 
     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.SERIALIZED_MODEL_FILE)
-        self.serialized_input_path = os.path.join(self.serialized_model_dir.name, settings.SERIALIZED_INPUT_FILE)
+        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.PICKLE_PROTOCOL)
-            torch.save(self.inputs, self.serialized_input_path, pickle_protocol=settings.PICKLE_PROTOCOL)
+            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
 
@@ -71,7 +71,7 @@ def _get_model_id(self) -> str:
         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.HASH_CHUNK_SIZE), b""):
+            for block in iter(lambda: serialized_model_file.read(settings.CENTML_HASH_CHUNK_SIZE), b""):
                 sha_hash.update(block)
 
         model_id = sha_hash.hexdigest()
@@ -80,7 +80,7 @@ def _get_model_id(self) -> str:
 
     def _download_model(self, model_id: str):
         download_response = requests.get(
-            url=f"{settings.CENTML_SERVER_URL}/download/{model_id}", timeout=settings.TIMEOUT
+            url=f"{settings.CENTML_SERVER_URL}/download/{model_id}", timeout=settings.CENTML_COMPILER_TIMEOUT
         )
         if download_response.status_code != HTTPStatus.OK:
             raise Exception(
@@ -106,7 +106,7 @@ def _compile_model(self, model_id: str):
             compile_response = requests.post(
                 url=f"{settings.CENTML_SERVER_URL}/submit/{model_id}",
                 files={"model": model_file, "inputs": input_file},
-                timeout=settings.TIMEOUT,
+                timeout=settings.CENTML_COMPILER_TIMEOUT,
             )
         if compile_response.status_code != HTTPStatus.OK:
             raise Exception(
@@ -120,7 +120,7 @@ def _wait_for_status(self, model_id: str) -> bool:
             status = None
             try:
                 status_response = requests.get(
-                    f"{settings.CENTML_SERVER_URL}/status/{model_id}", timeout=settings.TIMEOUT
+                    f"{settings.CENTML_SERVER_URL}/status/{model_id}", timeout=settings.CENTML_COMPILER_TIMEOUT
                 )
                 if status_response.status_code != HTTPStatus.OK:
                     raise Exception(
@@ -144,10 +144,10 @@ def _wait_for_status(self, model_id: str) -> bool:
             else:
                 tries += 1
 
-            if tries > settings.MAX_RETRIES:
+            if tries > settings.CENTML_COMPILER_MAX_RETRIES:
                 raise Exception("Waiting for status: compilation failed too many times.\n")
 
-            time.sleep(settings.COMPILING_SLEEP_TIME)
+            time.sleep(settings.CENTML_COMPILER_SLEEP_TIME)
 
     def remote_compilation_starter(self):
         try:

centml/compiler/config.py

@@ -10,26 +10,26 @@ class CompilationStatus(Enum):
 
 
 class Config(BaseSettings):
-    TIMEOUT: int = 10
-    MAX_RETRIES: int = 3
-    COMPILING_SLEEP_TIME: int = 15
+    CENTML_COMPILER_TIMEOUT: int = 10
+    CENTML_COMPILER_MAX_RETRIES: int = 3
+    CENTML_COMPILER_SLEEP_TIME: int = 15
 
-    CENTML_CACHE_DIR: str = os.path.expanduser("~/.cache/centml")
-    BACKEND_BASE_PATH: str = os.path.join(CENTML_CACHE_DIR, "backend")
-    SERVER_BASE_PATH: str = os.path.join(CENTML_CACHE_DIR, "server")
+    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.
-    SERIALIZED_MODEL_FILE: str = "serialized_model.zip"
-    SERIALIZED_INPUT_FILE: str = "serialized_input.zip"
-    PICKLE_PROTOCOL: int = 4
+    CENTML_SERIALIZED_MODEL_FILE: str = "serialized_model.zip"
+    CENTML_SERIALIZED_INPUT_FILE: str = "serialized_input.zip"
+    CENTML_PICKLE_PROTOCOL: int = 4
 
-    HASH_CHUNK_SIZE: int = 4096
+    CENTML_HASH_CHUNK_SIZE: int = 4096
 
     # If the server response is smaller than this, don't gzip it
-    MINIMUM_GZIP_SIZE: int = 1000
+    CENTML_MINIMUM_GZIP_SIZE: int = 1000
 
 
 settings = Config()

centml/compiler/server.py

@@ -14,11 +14,11 @@
 from centml.compiler.utils import get_server_compiled_forward_path
 
 app = FastAPI()
-app.add_middleware(GZipMiddleware, minimum_size=settings.MINIMUM_GZIP_SIZE)  # type: ignore
+app.add_middleware(GZipMiddleware, minimum_size=settings.CENTML_MINIMUM_GZIP_SIZE)  # type: ignore
 
 
 def get_status(model_id: str):
-    if not os.path.isdir(os.path.join(settings.SERVER_BASE_PATH, model_id)):
+    if not os.path.isdir(os.path.join(settings.CENTML_SERVER_BASE_PATH, model_id)):
         return CompilationStatus.NOT_FOUND
 
     if not os.path.isfile(get_server_compiled_forward_path(model_id)):
@@ -50,7 +50,7 @@ def background_compile(model_id: str, tfx_graph, example_inputs):
         # To avoid this, we write to a tmp file and rename it to the correct path after saving.
         save_path = get_server_compiled_forward_path(model_id)
         tmp_path = save_path + ".tmp"
-        torch.save(compiled_graph_module, tmp_path, pickle_protocol=settings.PICKLE_PROTOCOL)
+        torch.save(compiled_graph_module, tmp_path, pickle_protocol=settings.CENTML_PICKLE_PROTOCOL)
         os.rename(tmp_path, save_path)
     except Exception as e:
         logging.getLogger(__name__).exception(f"Saving graph module failed: {e}")
@@ -93,7 +93,7 @@ async def compile_model_handler(model_id: str, model: UploadFile, inputs: Upload
         return Response(status_code=200)
 
     # This effectively sets the model's status to COMPILING
-    os.makedirs(os.path.join(settings.SERVER_BASE_PATH, model_id))
+    os.makedirs(os.path.join(settings.CENTML_SERVER_BASE_PATH, model_id))
 
     tfx_graph, example_inputs = read_upload_files(model_id, model, inputs)
 

centml/compiler/utils.py

@@ -4,18 +4,18 @@
 
 
 def get_backend_compiled_forward_path(model_id: str):
-    os.makedirs(os.path.join(settings.BACKEND_BASE_PATH, model_id), exist_ok=True)
-    return os.path.join(settings.BACKEND_BASE_PATH, model_id, "compilation_return.pkl")
+    os.makedirs(os.path.join(settings.CENTML_BACKEND_BASE_PATH, model_id), exist_ok=True)
+    return os.path.join(settings.CENTML_BACKEND_BASE_PATH, model_id, "compilation_return.pkl")
 
 
 def get_server_compiled_forward_path(model_id: str):
-    os.makedirs(os.path.join(settings.SERVER_BASE_PATH, model_id), exist_ok=True)
-    return os.path.join(settings.SERVER_BASE_PATH, model_id, "compilation_return.pkl")
+    os.makedirs(os.path.join(settings.CENTML_SERVER_BASE_PATH, model_id), exist_ok=True)
+    return os.path.join(settings.CENTML_SERVER_BASE_PATH, model_id, "compilation_return.pkl")
 
 
 # This function will delete the storage_path/{model_id} directory
 def dir_cleanup(model_id: str):
-    dir_path = os.path.join(settings.SERVER_BASE_PATH, model_id)
+    dir_path = os.path.join(settings.CENTML_SERVER_BASE_PATH, model_id)
     if not os.path.exists(dir_path):
         return  # Directory does not exist, return
 

tests/test_backend.py

@@ -119,7 +119,7 @@ def test_successful_download(self, mock_requests, mock_load, mock_open, mock_mak
 
 
 class TestWaitForStatus(SetUpGraphModule):
-    @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
+    @patch("centml.compiler.config.settings.CENTML_COMPILER_SLEEP_TIME", new=0)
     @patch("centml.compiler.backend.requests")
     @patch("logging.Logger.exception")
     def test_invalid_status(self, mock_logger, mock_requests):
@@ -132,13 +132,13 @@ def test_invalid_status(self, mock_logger, mock_requests):
             self.runner._wait_for_status(model_id)
 
         mock_requests.get.assert_called()
-        assert mock_requests.get.call_count == settings.MAX_RETRIES + 1
-        assert len(mock_logger.call_args_list) == settings.MAX_RETRIES + 1
+        assert mock_requests.get.call_count == settings.CENTML_COMPILER_MAX_RETRIES + 1
+        assert len(mock_logger.call_args_list) == settings.CENTML_COMPILER_MAX_RETRIES + 1
         print(mock_logger.call_args_list)
         assert mock_logger.call_args_list[0].startswith("Status check failed:")
         assert "Waiting for status: compilation failed too many times.\n" == str(context.exception)
 
-    @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
+    @patch("centml.compiler.config.settings.CENTML_COMPILER_SLEEP_TIME", new=0)
     @patch("centml.compiler.backend.requests")
     @patch("logging.Logger.exception")
     def test_exception_in_status(self, mock_logger, mock_requests):
@@ -150,11 +150,11 @@ def test_exception_in_status(self, mock_logger, mock_requests):
             self.runner._wait_for_status(model_id)
 
         mock_requests.get.assert_called()
-        assert mock_requests.get.call_count == settings.MAX_RETRIES + 1
+        assert mock_requests.get.call_count == settings.CENTML_COMPILER_MAX_RETRIES + 1
         mock_logger.assert_called_with(f"Status check failed:\n{exception_message}")
         assert str(context.exception) == "Waiting for status: compilation failed too many times.\n"
 
-    @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
+    @patch("centml.compiler.config.settings.CENTML_COMPILER_SLEEP_TIME", new=0)
     @patch("centml.compiler.backend.Runner._compile_model")
     @patch("centml.compiler.backend.requests")
     def test_max_tries(self, mock_requests, mock_compile):
@@ -167,10 +167,10 @@ def test_max_tries(self, mock_requests, mock_compile):
         with self.assertRaises(Exception) as context:
             self.runner._wait_for_status(model_id)
 
-        self.assertEqual(mock_compile.call_count, settings.MAX_RETRIES + 1)
+        self.assertEqual(mock_compile.call_count, settings.CENTML_COMPILER_MAX_RETRIES + 1)
         self.assertIn("Waiting for status: compilation failed too many times", str(context.exception))
 
-    @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
+    @patch("centml.compiler.config.settings.CENTML_COMPILER_SLEEP_TIME", new=0)
     @patch("centml.compiler.backend.requests")
     def test_wait_on_compilation(self, mock_requests):
         # Mock the status check
@@ -186,7 +186,7 @@ def test_wait_on_compilation(self, mock_requests):
         # _wait_for_status should return True when compilation DONE
         assert self.runner._wait_for_status(model_id)
 
-    @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
+    @patch("centml.compiler.config.settings.CENTML_COMPILER_SLEEP_TIME", new=0)
     @patch("centml.compiler.backend.requests")
     @patch("centml.compiler.backend.Runner._compile_model")
     @patch("logging.Logger.exception")
@@ -206,10 +206,10 @@ def test_exception_in_compilation(self, mock_logger, mock_compile, mock_requests
             self.runner._wait_for_status(model_id)
 
         mock_requests.get.assert_called()
-        assert mock_requests.get.call_count == settings.MAX_RETRIES + 1
+        assert mock_requests.get.call_count == settings.CENTML_COMPILER_MAX_RETRIES + 1
 
         mock_compile.assert_called()
-        assert mock_compile.call_count == settings.MAX_RETRIES + 1
+        assert mock_compile.call_count == settings.CENTML_COMPILER_MAX_RETRIES + 1
 
         mock_logger.assert_called_with(f"Submitting compilation failed:\n{exception_message}")
         assert str(context.exception) == "Waiting for status: compilation failed too many times.\n"