From db5f5abd9e6cba1335e3a912eb4daaf77de44d04 Mon Sep 17 00:00:00 2001
From: Feroze Mohideen <feroze@porter.run>
Date: Fri, 17 May 2024 12:45:45 -0400
Subject: [PATCH] preflight check for model addons

---
 .../AddonFormContextProvider.tsx              | 104 +++++++++++++++++-
 .../home/add-on-dashboard/AddonTemplates.tsx  |  17 ++-
 2 files changed, 114 insertions(+), 7 deletions(-)

diff --git a/dashboard/src/main/home/add-on-dashboard/AddonFormContextProvider.tsx b/dashboard/src/main/home/add-on-dashboard/AddonFormContextProvider.tsx
index 4789d9165b..6e902e66f2 100644
--- a/dashboard/src/main/home/add-on-dashboard/AddonFormContextProvider.tsx
+++ b/dashboard/src/main/home/add-on-dashboard/AddonFormContextProvider.tsx
@@ -1,5 +1,6 @@
 import React, { createContext, useMemo, useState } from "react";
 import { zodResolver } from "@hookform/resolvers/zod";
+import { Contract } from "@porter-dev/api-contracts";
 import { useQueryClient } from "@tanstack/react-query";
 import { FormProvider, useForm } from "react-hook-form";
 import { useHistory } from "react-router";
@@ -8,11 +9,20 @@ import styled from "styled-components";
 import Loading from "components/Loading";
 import { Error as ErrorComponent } from "components/porter/Error";
 import { clientAddonValidator, type ClientAddon } from "lib/addons";
+import { updateExistingClusterContract } from "lib/clusters";
+import { type ClientPreflightCheck } from "lib/clusters/types";
 import { useAddon } from "lib/hooks/useAddon";
-import { getErrorMessageFromNetworkCall } from "lib/hooks/useCluster";
+import {
+  getErrorMessageFromNetworkCall,
+  preflightChecks,
+} from "lib/hooks/useCluster";
 import { useDefaultDeploymentTarget } from "lib/hooks/useDeploymentTarget";
 
-import { type UpdateClusterButtonProps } from "../infrastructure-dashboard/ClusterFormContextProvider";
+import { useClusterContext } from "../infrastructure-dashboard/ClusterContextProvider";
+import ClusterFormContextProvider, {
+  type UpdateClusterButtonProps,
+} from "../infrastructure-dashboard/ClusterFormContextProvider";
+import PreflightChecksModal from "../infrastructure-dashboard/modals/PreflightChecksModal";
 
 type AddonFormContextType = {
   updateAddonButtonProps: UpdateClusterButtonProps;
@@ -43,6 +53,12 @@ const AddonFormContextProvider: React.FC<AddonFormContextProviderProps> = ({
   children,
 }) => {
   const [updateAddonError, setUpdateAddonError] = useState<string>("");
+  const [failingPreflightChecks, setFailingPreflightChecks] = useState<
+    ClientPreflightCheck[]
+  >([]);
+  const [isCheckingQuotas, setIsCheckingQuotas] = useState<boolean>(false);
+  const { cluster } = useClusterContext();
+
   const { defaultDeploymentTarget } = useDefaultDeploymentTarget();
   const { updateAddon } = useAddon();
   const queryClient = useQueryClient();
@@ -57,12 +73,81 @@ const AddonFormContextProvider: React.FC<AddonFormContextProviderProps> = ({
     formState: { isSubmitting, errors },
   } = addonForm;
 
+  const modelAddonPreflightChecks = async (
+    addon: ClientAddon,
+    projectId: number
+  ): Promise<ClientPreflightCheck[] | undefined> => {
+    // TODO: figure out why data.template is undefined here. If it is defined, we can use data.template.isModelTemplate instead of hardcoding the type
+    if (
+      addon.config.type === "deepgram" &&
+      cluster.contract?.config &&
+      cluster.contract.config.cluster.cloudProvider === "AWS"
+    ) {
+      let clientContract = cluster.contract.config;
+      if (
+        !clientContract.cluster.config.nodeGroups.some(
+          (n) => n.nodeGroupType === "CUSTOM"
+        )
+      ) {
+        clientContract = {
+          ...clientContract,
+          cluster: {
+            ...clientContract.cluster,
+            config: {
+              ...clientContract.cluster.config,
+              nodeGroups: [
+                ...clientContract.cluster.config.nodeGroups,
+                {
+                  nodeGroupType: "CUSTOM",
+                  instanceType: "g4dn.xlarge",
+                  minInstances: 0,
+                  maxInstances: 1,
+                },
+              ],
+            },
+          },
+        };
+      }
+      const contract = Contract.fromJsonString(
+        atob(cluster.contract.base64_contract),
+        {
+          ignoreUnknownFields: true,
+        }
+      );
+      const contractCluster = contract.cluster;
+      if (contractCluster) {
+        const newContract = new Contract({
+          ...contract,
+          cluster: updateExistingClusterContract(
+            clientContract,
+            contractCluster
+          ),
+        });
+        setIsCheckingQuotas(true);
+        const preflightCheckResults = await preflightChecks(
+          newContract,
+          projectId
+        );
+        return preflightCheckResults;
+      }
+    }
+  };
+
   const onSubmit = handleSubmit(async (data) => {
     if (!projectId) {
       return;
     }
+    setFailingPreflightChecks([]);
     setUpdateAddonError("");
     try {
+      const preflightCheckResults = await modelAddonPreflightChecks(
+        data,
+        projectId
+      );
+      if (preflightCheckResults) {
+        setFailingPreflightChecks(preflightCheckResults);
+      }
+      setIsCheckingQuotas(false);
       await updateAddon({
         projectId,
         deploymentTargetId: defaultDeploymentTarget.id,
@@ -91,6 +176,9 @@ const AddonFormContextProvider: React.FC<AddonFormContextProviderProps> = ({
       props.status = "loading";
       props.isDisabled = true;
     }
+    if (isCheckingQuotas) {
+      props.loadingText = "Checking quotas...";
+    }
 
     if (updateAddonError) {
       props.status = (
@@ -103,7 +191,7 @@ const AddonFormContextProvider: React.FC<AddonFormContextProviderProps> = ({
     }
 
     return props;
-  }, [isSubmitting, errors, errors?.name?.value]);
+  }, [isSubmitting, errors, errors?.name?.value, isCheckingQuotas]);
 
   if (!projectId) {
     return <Loading />;
@@ -120,6 +208,16 @@ const AddonFormContextProvider: React.FC<AddonFormContextProviderProps> = ({
         <FormProvider {...addonForm}>
           <form onSubmit={onSubmit}>{children}</form>
         </FormProvider>
+        {failingPreflightChecks.length > 0 && (
+          <ClusterFormContextProvider projectId={projectId}>
+            <PreflightChecksModal
+              onClose={() => {
+                setFailingPreflightChecks([]);
+              }}
+              preflightChecks={failingPreflightChecks}
+            />
+          </ClusterFormContextProvider>
+        )}
       </Wrapper>
     </AddonFormContext.Provider>
   );
diff --git a/dashboard/src/main/home/add-on-dashboard/AddonTemplates.tsx b/dashboard/src/main/home/add-on-dashboard/AddonTemplates.tsx
index 056c0762d5..d64c99f870 100644
--- a/dashboard/src/main/home/add-on-dashboard/AddonTemplates.tsx
+++ b/dashboard/src/main/home/add-on-dashboard/AddonTemplates.tsx
@@ -17,6 +17,7 @@ import addOnGrad from "assets/add-on-grad.svg";
 import inferenceGrad from "assets/inference-grad.svg";
 
 import DashboardHeader from "../cluster-dashboard/DashboardHeader";
+import ClusterContextProvider from "../infrastructure-dashboard/ClusterContextProvider";
 import AddonForm from "./AddonForm";
 import AddonFormContextProvider from "./AddonFormContextProvider";
 
@@ -24,7 +25,7 @@ type Props = {
   filterModels?: boolean;
 };
 const AddonTemplates: React.FC<Props> = ({ filterModels }) => {
-  const { currentProject } = useContext(Context);
+  const { currentProject, currentCluster } = useContext(Context);
   const { search } = useLocation();
   const queryParams = new URLSearchParams(search);
   const history = useHistory();
@@ -40,9 +41,17 @@ const AddonTemplates: React.FC<Props> = ({ filterModels }) => {
 
   if (templateMatch) {
     return (
-      <AddonFormContextProvider projectId={currentProject?.id} redirectOnSubmit>
-        <AddonForm template={templateMatch} filterModels={filterModels} />
-      </AddonFormContextProvider>
+      <ClusterContextProvider
+        clusterId={currentCluster?.id}
+        refetchInterval={0}
+      >
+        <AddonFormContextProvider
+          projectId={currentProject?.id}
+          redirectOnSubmit
+        >
+          <AddonForm template={templateMatch} filterModels={filterModels} />
+        </AddonFormContextProvider>
+      </ClusterContextProvider>
     );
   }