updates stack handling

fuzzylabs · Aug 15, 2023 · b42b31d · b42b31d
1 parent 407be1c
commit b42b31d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/src/matcha_ml/core/core.py b/src/matcha_ml/core/core.py
@@ -18,7 +18,7 @@
 from matcha_ml.services.global_parameters_service import GlobalParameters
 from matcha_ml.state import MatchaStateService, RemoteStateManager
 from matcha_ml.state.matcha_state import MatchaState
-from matcha_ml.templates.azure_template import AzureTemplate
+from matcha_ml.templates.azure_template import DEFAULT_STACK, LLM_STACK, AzureTemplate
 
 
 class StackTypeMeta(
@@ -293,8 +293,12 @@ def provision(
             "infrastructure",
             stack_name,
         )
+        if stack_name == "llm":
+            submodule_names = DEFAULT_STACK + LLM_STACK
+        else:
+            submodule_names = DEFAULT_STACK
 
-        azure_template = AzureTemplate()
+        azure_template = AzureTemplate(submodule_names)
 
         zenml_version = infer_zenml_version()
         config = azure_template.build_template_configuration(

diff --git a/src/matcha_ml/templates/azure_template.py b/src/matcha_ml/templates/azure_template.py
@@ -1,10 +1,10 @@
 """Build a template for provisioning resources on Azure using terraform files."""
-from typing import Optional
+from typing import List, Optional
 
 from matcha_ml.state import MatchaState, MatchaStateService
 from matcha_ml.templates.base_template import BaseTemplate, TemplateVariables
 
-SUBMODULE_NAMES = [
+DEFAULT_STACK = [
     "aks",
     "resource_group",
     "mlflow_module",
@@ -16,8 +16,8 @@
     "zen_server/zenml_helm",
     "zen_server/zenml_helm/templates",
     "data_version_control_storage",
-    "chroma",
 ]
+LLM_STACK = ["chroma"]
 
 
 class AzureTemplate(BaseTemplate):
@@ -27,13 +27,13 @@ class AzureTemplate(BaseTemplate):
         BaseTemplate: The base template class.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, submodule_names: List[str]) -> None:
         """Initialize the StateStorageTemplate with the submodule names.
 
         Args:
             submodule_names (List[str]): A list of submodule names.
         """
-        super().__init__(SUBMODULE_NAMES)
+        super().__init__(submodule_names)
 
     def build_template(
         self,