From e5d5f4be105822263ce43518e3bcd4c3f758b4f0 Mon Sep 17 00:00:00 2001 From: Ryan Weiler Date: Mon, 26 Oct 2026 16:55:09 -0400 Subject: [PATCH] feat: removing old kserve implementation --- docs/python_gaas_tools/gaas_model.md | 4 +- .../app_root/version/assets/public/SKILL.md | 7 +- py/gaas_gpt_model.py | 38 - src/prerna/cluster/util/IModelZKListener.java | 53 - .../cluster/util/IRemoteClientServer.java | 110 -- .../cluster/util/ModelAvailableListener.java | 98 -- .../cluster/util/ModelFailListener.java | 106 -- .../cluster/util/ModelInitListener.java | 94 -- .../util/ModelUnavailableListener.java | 88 -- src/prerna/cluster/util/ModelZKServer.java | 773 ---------- .../cluster/util/RemoteClientServerZK.java | 565 ------- .../util/RemoteClientServerZKRESTProxy.java | 609 -------- src/prerna/cluster/util/RemoteModelInfo.java | 50 - src/prerna/cluster/util/ZKClientFactory.java | 64 - src/prerna/engine/api/ModelTypeEnum.java | 10 - .../engine/api/RemoteModelStateEnum.java | 36 - .../OpenAITranscribeFunctionEngine.java | 54 - .../impl/model/AbstractRemoteModelEngine.java | 539 ------- .../impl/model/KServeImageEmbedEngine.java | 81 - .../engine/impl/model/KServeImageEngine.java | 110 -- .../engine/impl/model/KServeTTSEngine.java | 104 -- .../engine/impl/model/KServeVisionEngine.java | 87 -- .../impl/model/KubernetesModelScaler.java | 305 ---- src/prerna/engine/impl/model/NEREngine.java | 95 -- .../impl/model/kserve/KServeAdapter.java | 142 -- .../AskImageModelEngineResponse.java | 86 +- .../AskStringModelEngineResponse.java | 41 - .../responses/AskTTSModelEngineResponse.java | 278 ---- .../responses/NerModelEngineResponse.java | 152 -- .../automation/AutomationConstants.java | 1 - .../AutomationDefinitionValidator.java | 5 - .../definition/AutomationNodeCatalog.java | 7 - .../definition/AutomationNodeType.java | 2 - .../AutomationPixelExpressionRenderer.java | 13 - .../run/AutomationRunExecutionService.java | 2 +- src/prerna/reactor/model/CanItRunReactor.java | 69 - .../model/GetNodePoolsInfoReactor.java | 63 - .../GetRemoteModelDeployConfigsReactor.java | 88 -- .../reactor/model/MyRemoteModelsStatus.java | 80 - src/prerna/reactor/model/NERReactor.java | 117 -- .../model/RemoteModelShutdownReactor.java | 73 - .../model/RemoteModelStartReactor.java | 83 -- src/prerna/sablecc2/om/ReactorKeysEnum.java | 3 - src/prerna/util/Settings.java | 5 +- .../util/ModelAvailableListenerUnitTests.java | 173 --- .../util/ModelFailListenerUnitTests.java | 159 -- .../util/ModelInitListenerUnitTests.java | 181 --- .../ModelUnavailableListenerUnitTests.java | 141 -- .../cluster/util/ModelZKServerUnitTests.java | 379 ----- ...emoteClientServerZKRESTProxyUnitTests.java | 1304 ----------------- .../util/RemoteClientServerZKUnitTests.java | 839 ----------- .../util/RemoteModelInfoUnitTests.java | 145 -- .../util/ZKClientFactoryUnitTests.java | 77 - .../engine/api/ModelTypeEnumUnitTests.java | 24 - .../api/RemoteModelStateEnumUnitTests.java | 65 - .../model/kserve/KServeAdapterUnitTests.java | 126 -- .../AskImageModelEngineResponseUnitTests.java | 93 -- .../NerModelEngineResponseUnitTests.java | 107 -- ...ationPixelExpressionRendererUnitTests.java | 6 +- 59 files changed, 8 insertions(+), 9201 deletions(-) delete mode 100644 src/prerna/cluster/util/IModelZKListener.java delete mode 100644 src/prerna/cluster/util/IRemoteClientServer.java delete mode 100644 src/prerna/cluster/util/ModelAvailableListener.java delete mode 100644 src/prerna/cluster/util/ModelFailListener.java delete mode 100644 src/prerna/cluster/util/ModelInitListener.java delete mode 100644 src/prerna/cluster/util/ModelUnavailableListener.java delete mode 100644 src/prerna/cluster/util/ModelZKServer.java delete mode 100644 src/prerna/cluster/util/RemoteClientServerZK.java delete mode 100644 src/prerna/cluster/util/RemoteClientServerZKRESTProxy.java delete mode 100644 src/prerna/cluster/util/RemoteModelInfo.java delete mode 100644 src/prerna/cluster/util/ZKClientFactory.java delete mode 100644 src/prerna/engine/api/RemoteModelStateEnum.java delete mode 100644 src/prerna/engine/impl/model/AbstractRemoteModelEngine.java delete mode 100644 src/prerna/engine/impl/model/KServeImageEmbedEngine.java delete mode 100644 src/prerna/engine/impl/model/KServeImageEngine.java delete mode 100644 src/prerna/engine/impl/model/KServeTTSEngine.java delete mode 100644 src/prerna/engine/impl/model/KServeVisionEngine.java delete mode 100644 src/prerna/engine/impl/model/KubernetesModelScaler.java delete mode 100644 src/prerna/engine/impl/model/NEREngine.java delete mode 100644 src/prerna/engine/impl/model/kserve/KServeAdapter.java delete mode 100644 src/prerna/engine/impl/model/responses/AskTTSModelEngineResponse.java delete mode 100644 src/prerna/engine/impl/model/responses/NerModelEngineResponse.java delete mode 100644 src/prerna/reactor/model/CanItRunReactor.java delete mode 100644 src/prerna/reactor/model/GetNodePoolsInfoReactor.java delete mode 100644 src/prerna/reactor/model/GetRemoteModelDeployConfigsReactor.java delete mode 100644 src/prerna/reactor/model/MyRemoteModelsStatus.java delete mode 100644 src/prerna/reactor/model/NERReactor.java delete mode 100644 src/prerna/reactor/model/RemoteModelShutdownReactor.java delete mode 100644 src/prerna/reactor/model/RemoteModelStartReactor.java delete mode 100644 test/prerna/cluster/util/ModelAvailableListenerUnitTests.java delete mode 100644 test/prerna/cluster/util/ModelFailListenerUnitTests.java delete mode 100644 test/prerna/cluster/util/ModelInitListenerUnitTests.java delete mode 100644 test/prerna/cluster/util/ModelUnavailableListenerUnitTests.java delete mode 100644 test/prerna/cluster/util/ModelZKServerUnitTests.java delete mode 100644 test/prerna/cluster/util/RemoteClientServerZKRESTProxyUnitTests.java delete mode 100644 test/prerna/cluster/util/RemoteClientServerZKUnitTests.java delete mode 100644 test/prerna/cluster/util/RemoteModelInfoUnitTests.java delete mode 100644 test/prerna/cluster/util/ZKClientFactoryUnitTests.java delete mode 100644 test/prerna/engine/api/RemoteModelStateEnumUnitTests.java delete mode 100644 test/prerna/engine/impl/model/kserve/KServeAdapterUnitTests.java delete mode 100644 test/prerna/engine/impl/model/responses/NerModelEngineResponseUnitTests.java diff --git a/docs/python_gaas_tools/gaas_model.md b/docs/python_gaas_tools/gaas_model.md index 7ea14579caa..b7f352d0932 100644 --- a/docs/python_gaas_tools/gaas_model.md +++ b/docs/python_gaas_tools/gaas_model.md @@ -27,8 +27,6 @@ The `py/gaas_gpt_model.py` module provides classes for interacting with various * Constructs a Pixel script: `LLM(engine="", command="", useHistory=, context=[""], paramValues=[]);`. * Executes this Pixel script via `super().callReactor()`. * Returns the response from the LLM, typically including the generated text and token counts. - * `ner(text: str, entities: List[str], mask_entities: List[str] = [], param_dict: Optional[Dict] = None, insight_id: Optional[str] = None)`: - * Executes an `NER(...)` Pixel command for Named Entity Recognition. * `get_conversation_history(insight_id: Optional[str] = None) -> List[Dict]`: * Executes a `GetRoomMessages(roomId="");` Pixel script to retrieve chat history from the `ModelInferenceLogsDatabase`. * `embeddings(strings_to_embed: List[str], param_dict: Optional[Dict] = None, insight_id: Optional[str] = None) -> List[Dict]`: @@ -74,7 +72,7 @@ The `py/gaas_gpt_model.py` module provides classes for interacting with various * **Initialization `__init__(self, model_engine_class: Optional[str] = "TOMCAT", **kwargs)`**: * `model_engine_class` (str, default: "TOMCAT"): Determines the type of engine to create. Can be "TOMCAT" or "LOCAL". "HF_PIPELINE" is mentioned but not fully implemented in the constructor logic shown. * `**kwargs`: Passed to the constructor of the chosen model engine class (e.g., `engine_id` for `TomcatModelEngine`). -* **Key Methods**: All methods (`get_model_type`, `ask`, `instruct`, `embeddings`, `keyword_extraction`, `ner`, `do_call`, `get_model_engine_id`, `get_conversation_history`) are wrappers that delegate the call to the underlying `self.model_engine` instance. +* **Key Methods**: All methods (`get_model_type`, `ask`, `instruct`, `embeddings`, `keyword_extraction`, `do_call`, `get_model_engine_id`, `get_conversation_history`) are wrappers that delegate the call to the underlying `self.model_engine` instance. * **Langchain Integration**: * `to_langchain_embedder()`: Wraps the `ModelEngine` to make its `embeddings` method compatible with Langchain's `Embeddings` interface. * `to_langchain_chat_model()`: Wraps the `ModelEngine` to make its `ask` and `get_conversation_history` methods compatible with Langchain's `BaseChatModel` interface. This includes converting message formats. diff --git a/project/platform__model/app_root/version/assets/public/SKILL.md b/project/platform__model/app_root/version/assets/public/SKILL.md index a2392a0fc0f..6b810ac86ba 100644 --- a/project/platform__model/app_root/version/assets/public/SKILL.md +++ b/project/platform__model/app_root/version/assets/public/SKILL.md @@ -1,6 +1,6 @@ --- name: model -description: Use when writing code in an app that calls an LLM, embedding model, or other model engine, OR when listing/selecting models the user has access to. Covers the LLM() and MyEngines() pixel commands via @semoss/sdk's runPixel, including prompt/completion calls, conversational history, structured outputs, attaching media of any type (image, PDF, Word, spreadsheet, audio, video) with the media argument, and parsing model responses. Also covers which top-level arguments full_prompt drops, media included, plus the other model pixels: Embeddings(), MultiModalEmbeddings(), Vision(), NER(), GetContextWindow(), GetUserModelUsage(), and the provider batch API (BatchLLM, GetModelBatchStatus, GetModelBatchResults). Do not use for vector database queries (see vector). +description: Use when writing code in an app that calls an LLM, embedding model, or other model engine, OR when listing/selecting models the user has access to. Covers the LLM() and MyEngines() pixel commands via @semoss/sdk's runPixel, including prompt/completion calls, conversational history, structured outputs, attaching media of any type (image, PDF, Word, spreadsheet, audio, video) with the media argument, and parsing model responses. Also covers which top-level arguments full_prompt drops, media included, plus the other model pixels: Embeddings(), MultiModalEmbeddings(), Vision(), GetContextWindow(), GetUserModelUsage(), and the provider batch API (BatchLLM, GetModelBatchStatus, GetModelBatchResults). Do not use for vector database queries (see vector). --- # Model Engine @@ -212,14 +212,13 @@ Embeddings(engine="${EMBEDDING_MODEL_ID}", values=["first text", "second text"]) Returns a map with the vectors in `pixelReturn[0].output`. If the engine is not an embedding model, the output is the literal string "This model does not support embeddings." - check for it. `MultiModalEmbeddings(engine=, text=, image=, video=)` embeds any combination of modalities (image/video accept base64, data URLs, or remote URLs), with results broken out by modality. -### Vision and NER +### Vision ``` Vision(engine="${MODEL_ID}", command=["Describe this diagram"], image=["${urlOrBase64}"]); -NER(engine="${NER_MODEL_ID}", prompt=["${text}"], entities=["PERSON","EMAIL"], maskEntities=[true]); ``` -`Vision` is the image-text-to-text task for models exposing it. `NER` extracts (and with `maskEntities` redacts) named entities - the engine must actually be an NER model or the call fails. +`Vision` is the image-text-to-text task for models exposing it. ### Token budget and usage diff --git a/py/gaas_gpt_model.py b/py/gaas_gpt_model.py index 3b62ccfe4e8..45be523e6e8 100644 --- a/py/gaas_gpt_model.py +++ b/py/gaas_gpt_model.py @@ -206,35 +206,6 @@ def ask( return pixelReturn - def ner( - self, - text: str, - entities: List[str], - mask_entities: List[str] = [], - param_dict: Optional[Dict] = None, - insight_id: Optional[str] = None, - ): - if insight_id is None: - insight_id = self.insight_id - - epoc = super().get_next_epoc() - - pixel = f'NER(engine="{self.engine_id}", prompt="{text}", entities={entities}, maskEntities={mask_entities});' - - pixelReturn = super().callReactor( - epoc=epoc, - pixel=pixel, - insight_id=insight_id, - ) - - if pixelReturn is not None and len(pixelReturn) > 0: - output = pixelReturn[0]["pixelReturn"][0] - if "ERROR" in output["operationType"]: - raise Exception(output["output"]) - return output["output"] - - return pixelReturn - def get_conversation_history(self, insight_id: Optional[str] = None) -> List[Dict]: """This method is responsible to get message history back from the model logs database. @@ -608,15 +579,6 @@ def keyword_extraction( ): return self.model_engine.keyword_extraction(**kwargs) - def ner( - self, - insight_id: Optional[ - str - ] = None, # TODO remove once users stop using it. No longer needs to be set. - **kwargs, - ): - return self.model_engine.ner(**kwargs) - def do_call(self, **kwargs: Any) -> Any: return self.model_engine.embeddings(**kwargs) diff --git a/src/prerna/cluster/util/IModelZKListener.java b/src/prerna/cluster/util/IModelZKListener.java deleted file mode 100644 index ac76a04b7d5..00000000000 --- a/src/prerna/cluster/util/IModelZKListener.java +++ /dev/null @@ -1,53 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.util.List; - -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.ZooKeeper; - -public interface IModelZKListener { - - // needs the zookeeper ? - // need the model zk server this is where it puts everything - // needs the list of events as an array or a list - // needs path - // needs the predicate - starts with, ends with, contains - - // finally a process - public void setModelZK(ModelZKServer server); - - public List getEvents(); - - public String getPath(); - - public List getPredicates(); - - public void process(String path, ZooKeeper zk); -} diff --git a/src/prerna/cluster/util/IRemoteClientServer.java b/src/prerna/cluster/util/IRemoteClientServer.java deleted file mode 100644 index f1ba113d843..00000000000 --- a/src/prerna/cluster/util/IRemoteClientServer.java +++ /dev/null @@ -1,110 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.util.List; -import java.util.Map; - -import prerna.engine.api.RemoteModelStateEnum; - -/** - * Interface to abstract the remote client server functionality. - * This allows both ZooKeeper-based and REST-based implementations. - */ -public interface IRemoteClientServer { - - /** - * Gets the current state of a model - * @param modelId The model ID - * @return The current state of the model - */ - RemoteModelStateEnum getModelState(String modelId); - - /** - * Gets the cluster IP address for the specified model - * @param modelId The model ID - * @return The cluster IP address, or null if not found - */ - String getModelClusterIp(String modelId); - - /** - * Gets the cluster IP address for the KMS server - * @return The cluster IP address, or null if not found - */ - String getModelScalerIp(); - - /** - * Gets the friendly name of the model - * @param modelId The model ID - * @return The model name, or null if not found - */ - String getModelName(String modelId); - - /** - * Checks if the model is currently warming up - * @param modelId The model ID - * @return true if the model is warming, false otherwise - */ - boolean isModelWarming(String modelId); - - /** - * Checks if the model is currently active - * @param modelId The model ID - * @return true if the model is active, false otherwise - */ - boolean isModelActive(String modelId); - - /** - * Waits for a model to become active - * @param modelId The model ID - * @param timeoutMs Maximum time to wait in milliseconds - * @return true if the model became active within the timeout, false otherwise - */ - boolean waitForModelActive(String modelId, long timeoutMs); - - /** - * Waits for a model to reach a specific state - * @param modelId The model ID - * @param desiredState The state to wait for - * @param timeoutMs Maximum time to wait in milliseconds - * @return true if the model reached the desired state within the timeout, false otherwise - */ - boolean waitForState(String modelId, RemoteModelStateEnum desiredState, long timeoutMs); - - /** - * Gets a list of active models - * @return List of RemoteModelInfo objects for active models - */ - List getActiveModels(); - - /** - * Gets a list of warming models - * @return List of RemoteModelInfo objects for warming models - */ - List getWarmingModels(); -} \ No newline at end of file diff --git a/src/prerna/cluster/util/ModelAvailableListener.java b/src/prerna/cluster/util/ModelAvailableListener.java deleted file mode 100644 index 33c032593f1..00000000000 --- a/src/prerna/cluster/util/ModelAvailableListener.java +++ /dev/null @@ -1,98 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.io.UnsupportedEncodingException; -import java.util.ArrayList; -import java.util.List; - -import org.apache.zookeeper.KeeperException; -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.data.Stat; -import org.apache.zookeeper.ZooKeeper; - -public class ModelAvailableListener implements IModelZKListener { - - String path = null; - public static final String AVAILABLE = "Available"; - ModelZKServer server = null; - - public ModelAvailableListener(String path, ModelZKServer server) - { - this.server = server; - this.path = path; - } - - @Override - public void setModelZK(ModelZKServer server) { - // TODO Auto-generated method stub - this.server = server; - } - - @Override - public List getEvents() { - // TODO Auto-generated method stub - List retList = new ArrayList(); - retList.add(EventType.NodeDataChanged); - return retList; - } - - @Override - public String getPath() { - // TODO Auto-generated method stub - return this.path; - } - - @Override - public List getPredicates() { - // TODO Auto-generated method stub - List predicates = new ArrayList(); - predicates.add("equals"); - return predicates; - } - - @Override - public void process(String path, ZooKeeper zk) { - // TODO Auto-generated method stub - String status = server.getNodeData(path); - if(status.equalsIgnoreCase("Available")) - { - System.err.println("We are done here.. "); - String endpointNode = path.replace("status", "endpoint"); - String endpoint = server.getNodeData(endpointNode); - - System.err.println("Endpoint " + endpoint); - } - // need to account if this is a fail - if(status.equalsIgnoreCase("Fail")) - { - String endpointNode = path.replace("status", "endpoint"); - System.err.println("Could not satisfy request"); - } - } -} diff --git a/src/prerna/cluster/util/ModelFailListener.java b/src/prerna/cluster/util/ModelFailListener.java deleted file mode 100644 index 49f5905b378..00000000000 --- a/src/prerna/cluster/util/ModelFailListener.java +++ /dev/null @@ -1,106 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.io.UnsupportedEncodingException; - -import java.util.ArrayList; -import java.util.List; - -import org.apache.zookeeper.KeeperException; -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.data.Stat; -import org.apache.zookeeper.ZooKeeper; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import prerna.util.Constants; - -// counts the number of fails -// and says if the overall status failed - -public class ModelFailListener implements IModelZKListener { - - private static final Logger classLogger = LogManager.getLogger(ModelFailListener.class); - - String path = null; - ModelZKServer server = null; - - public ModelFailListener(String path, ModelZKServer server) - { - this.server = server; - this.path = path; - } - - @Override - public void setModelZK(ModelZKServer server) { - // TODO Auto-generated method stub - this.server = server; - } - - @Override - public List getEvents() { - // TODO Auto-generated method stub - List retList = new ArrayList(); - retList.add(EventType.NodeChildrenChanged); - return retList; - } - - @Override - public String getPath() { - // TODO Auto-generated method stub - return this.path; - } - - @Override - public List getPredicates() { - // TODO Auto-generated method stub - List predicates = new ArrayList(); - predicates.add("equals"); - return predicates; - } - - @Override - public void process(String path, ZooKeeper zk) { - // TODO Auto-generated method stub - try { - List failedChildren = zk.getChildren(path, true); - List serverChildren = zk.getChildren("/server", true); - - String statusNode = path.replace("fail", "status"); - - if(failedChildren.size() == serverChildren.size()) - server.updateNodeData(statusNode, "FAIL", true); - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - } -} diff --git a/src/prerna/cluster/util/ModelInitListener.java b/src/prerna/cluster/util/ModelInitListener.java deleted file mode 100644 index 2cdc3e147a9..00000000000 --- a/src/prerna/cluster/util/ModelInitListener.java +++ /dev/null @@ -1,94 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.io.UnsupportedEncodingException; - -import java.util.ArrayList; -import java.util.List; - -import org.apache.zookeeper.KeeperException; -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.data.Stat; -import org.apache.zookeeper.ZooKeeper; - -// counts the number of fails -// and says if the overall status failed - -public class ModelInitListener implements IModelZKListener { - - String path = null; - ModelZKServer server = null; - - public ModelInitListener(String path, ModelZKServer server) - { - this.server = server; - this.path = path; - } - - @Override - public void setModelZK(ModelZKServer server) { - // TODO Auto-generated method stub - this.server = server; - } - - @Override - public List getEvents() { - // TODO Auto-generated method stub - List retList = new ArrayList(); - retList.add(EventType.NodeDataChanged); - return retList; - } - - @Override - public String getPath() { - // TODO Auto-generated method stub - return this.path; - } - - @Override - public List getPredicates() { - // TODO Auto-generated method stub - List predicates = new ArrayList(); - predicates.add("equals"); - return predicates; - } - - @Override - public void process(String path, ZooKeeper zk) { - // TODO Auto-generated method stub - String status = server.getNodeData(path); - String model = path.replace("/status", ""); - - String smss = server.getNodeData(model); - - if(status.equalsIgnoreCase("INIT") && !server.catchup) - server.spinModel(smss); - - } -} diff --git a/src/prerna/cluster/util/ModelUnavailableListener.java b/src/prerna/cluster/util/ModelUnavailableListener.java deleted file mode 100644 index d76eaa7fead..00000000000 --- a/src/prerna/cluster/util/ModelUnavailableListener.java +++ /dev/null @@ -1,88 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.io.UnsupportedEncodingException; - -import java.util.ArrayList; -import java.util.List; - -import org.apache.zookeeper.KeeperException; -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.data.Stat; -import org.apache.zookeeper.ZooKeeper; - -// happens when the endpoint disappears -// it will update the status back to INIT - -public class ModelUnavailableListener implements IModelZKListener { - - String path = null; - ModelZKServer server = null; - - public ModelUnavailableListener(String path, ModelZKServer server) - { - this.server = server; - this.path = path; - } - - @Override - public void setModelZK(ModelZKServer server) { - // TODO Auto-generated method stub - this.server = server; - } - - @Override - public List getEvents() { - // TODO Auto-generated method stub - List retList = new ArrayList(); - retList.add(EventType.NodeDeleted); - return retList; - } - - @Override - public String getPath() { - // TODO Auto-generated method stub - return this.path; - } - - @Override - public List getPredicates() { - // TODO Auto-generated method stub - List predicates = new ArrayList(); - predicates.add("equals"); - return predicates; - } - - @Override - public void process(String path, ZooKeeper zk) { - // TODO Auto-generated method stub - String endpointNode = path.replace("endpoint", "status"); - server.updateNodeData(endpointNode, "INIT", true); - } -} diff --git a/src/prerna/cluster/util/ModelZKServer.java b/src/prerna/cluster/util/ModelZKServer.java deleted file mode 100644 index 284071e2b19..00000000000 --- a/src/prerna/cluster/util/ModelZKServer.java +++ /dev/null @@ -1,773 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.io.BufferedReader; -import java.io.File; -import java.io.IOException; -import java.io.InputStreamReader; -import java.io.StringReader; -import java.io.UnsupportedEncodingException; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.Iterator; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import java.util.Vector; -import java.util.concurrent.TimeUnit; - -import org.apache.curator.framework.CuratorFramework; -import org.apache.curator.framework.CuratorFrameworkFactory; -import org.apache.curator.framework.recipes.cache.ChildData; -import org.apache.curator.framework.recipes.cache.CuratorCache; -import org.apache.curator.framework.recipes.cache.CuratorCacheListener; -import org.apache.curator.framework.recipes.locks.InterProcessMutex; -import org.apache.curator.retry.RetryOneTime; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.apache.zookeeper.AddWatchMode; -import org.apache.zookeeper.CreateMode; -import org.apache.zookeeper.KeeperException; -import org.apache.zookeeper.WatchedEvent; -import org.apache.zookeeper.Watcher; -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.ZooDefs.Ids; -import org.apache.zookeeper.ZooKeeper; -import org.apache.zookeeper.data.Stat; - -import com.google.gson.Gson; -import com.google.gson.JsonSyntaxException; - -import prerna.ds.py.PyTranslator; -import prerna.om.ClientProcessWrapper; -import prerna.om.Insight; -import prerna.om.InsightStore; -import prerna.tcp.client.NativePySocketClient; -import prerna.util.Constants; -import prerna.util.PortAllocator; -import prerna.util.Settings; -import prerna.util.Utility; - -public class ModelZKServer implements Watcher, CuratorCacheListener { - - /* - * // Environment Variables this depnds on zk - semicolon separated list of zk - * to use for registration home - what is the main registration root - default - * is assumed as /semoss_root app - The root for various apps - default is - * assumed as /app host - ip and port - this is useful if you are running - * multiple containers on the same box bu - the boot user who is being loaded on - * this box - defaults to generic DNS container address - * - * // registers to semoss_root/semoss- as a ephemeral_sequential // the - * semoss_root is an environment variable - this will allow me to spin as many - * cluster as I want Major version minor version semoss_sequential node host ip - * and port number of cpus memory User who is booted on it The registration is - * typically register_root / #cpu / #memory / semoss-node register_root / - * semoss-node nginx will only watch the resigter_root/semoss_node initially I - * could even do this for specific databases so the same thing is not loaded - * multiple times // zk e - * - */ - - public static final String ZK_SERVER = "zk"; - public static final String HOST = "host"; - public static final String TIMEOUT = "to"; - public static final String BOOTUSER = "bu"; - public static final String HOME = "zk_home"; - public static final String APP_HOME = "app"; - public static final String SEMOSS_HOME = "sem"; - - public static final String AVAILABLE = "available"; - public static final String MODEL_ROOT = "/model"; - public static final String SERVER_PATH = "/server"; - String workingDirectoryBasePath = null; - NativePySocketClient socketClient = null; - PyTranslator pyTranslator = null; - Process process = null; - String workingDirectory = null; - String prefix = null; - File cacheFolder = null; - private static final Logger classLogger = LogManager.getLogger(ModelZKServer.class); - Properties prop = null; - Gson gson = new Gson(); - Map modelLock = new HashMap(); - Map modelClient = new HashMap(); - List existingModels = new ArrayList(); - List supportedModels = new ArrayList(); - Map modelSMSS = new HashMap(); - String id = "RANDOM_ID"; - boolean catchup = true; - - public ZooKeeper zk = null; - Map env = null; - public String zkServer = "localhost:2181"; - public String host = "localhost:8888"; - public String user = "generic"; - public String home = "/semoss_root"; - public String container = "/container"; - public String app = "/app"; - public static String semossHome = "/opt/semosshome/"; - boolean init = false; - Map> listeners = new HashMap>(); - - boolean connected = false; - - public static ModelZKServer zkClient = null; - - Map repeat = new HashMap(); - - CuratorFramework client = null; - - int version = 0; - - protected ModelZKServer() { - - } - - public static ModelZKServer getInstance(String id) { - if (zkClient == null) { - zkClient = new ModelZKServer(); - zkClient.id = id; - zkClient.init(); - zkClient.initCurator(); - } - if (zkClient.connected) { - return zkClient; - } - return null; - } - - public void reconnect() { - zkClient.init(); - } - - public void initCurator() { - // make the curator here also - try { - client = CuratorFrameworkFactory.newClient(zkServer, new RetryOneTime(1)); - client.start(); - } catch (Exception e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - - } - - public void waitHere() { - try { - - BufferedReader br = new BufferedReader(new InputStreamReader(System.in)); - br.readLine(); - - } catch (Exception ignored) { - - } - } - - public void init() { - // initiates connection to zk and makes the connection - try { - env = System.getenv(); - if (env.containsKey(ZK_SERVER)) { - zkServer = env.get(ZK_SERVER); - } - - if (env.containsKey(ZK_SERVER.toUpperCase())) { - zkServer = env.get(ZK_SERVER.toUpperCase()); - } - - if (env.containsKey(HOST)) { - host = env.get(HOST); - } - - if (env.containsKey(HOST.toUpperCase())) { - host = env.get(HOST.toUpperCase()); - } - - int timeout = (30 * 60 * 1000); - - if (env.containsKey(TIMEOUT)) { - host = env.get(TIMEOUT); - } - - if (env.containsKey(TIMEOUT.toUpperCase())) { - host = env.get(TIMEOUT.toUpperCase()); - } - - if (env.containsKey(BOOTUSER)) { - user = env.get(BOOTUSER); - } - - if (env.containsKey(BOOTUSER.toUpperCase())) { - user = env.get(BOOTUSER.toUpperCase()); - } - - if (env.containsKey(HOME)) { - home = env.get(HOME); - } - - if (env.containsKey(HOME.toUpperCase())) { - home = env.get(HOME.toUpperCase()); - } - - if (env.containsKey(APP_HOME)) { - app = env.get(APP_HOME); - } - - if (env.containsKey(APP_HOME.toUpperCase())) { - app = env.get(APP_HOME.toUpperCase()); - } - - if (env.containsKey(SEMOSS_HOME)) { - semossHome = env.get(SEMOSS_HOME); - } - - if (env.containsKey(SEMOSS_HOME.toUpperCase())) { - semossHome = env.get(SEMOSS_HOME.toUpperCase()); - } - - // TODO >>>timb:not sure if the host is needed for both the engine and user - // containers - if (zkServer != null && host != null) { - // open zk - // default time is 30 min - zk = new ZooKeeper(zkServer, timeout, this); - connected = true; - initCurator(); - } - - // start a base python - - socketClient = connect2Py(null); - Insight processInsight = new Insight(); - InsightStore.getInstance().put(processInsight); - this.pyTranslator = new PyTranslator(socketClient, processInsight); - - // publish this node ? - do we even need to ? - addServer(); - - // catchup - catchup(); - - // listen for getChildren on /model - addCuratorListener(this, MODEL_ROOT); - - } catch (IOException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - } - - public void addServer() { - // this merely adds the server - - // create server if one does not exist - // gets the lock - // updates the server count - try { - Stat s = zk.exists(SERVER_PATH, false); - - if (s == null) { - String data = "1"; - zk.create(SERVER_PATH, data.getBytes(), Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); - } - // add this as an ephemeral sequential node - // although I dont know why I need sequential.. I could just add it - zk.create(SERVER_PATH + "/" + id, id.getBytes(), Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL); - } catch (NumberFormatException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (Exception e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - } - - @Override - public void process(WatchedEvent event) { - // TODO Auto-generated method stub - // System.out.println("Ok.. something came back on this" + arg0); - processEvent(event.getPath(), event); - } - - public String getNodeData(String path) { - try { - byte[] nodeBytes = zk.getData(path, true, new Stat()); - String nodeData = new String(nodeBytes, "UTF-8"); - return nodeData; - } catch (UnsupportedEncodingException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - return null; - } - - public void updateNodeData(String path, String data, boolean create) { - try { - Stat stat = zk.exists(path, false); - if (stat == null && create) { - Stat ver = new Stat(1, 1, 0, 0, 1, 0, 0, 1, data.length(), 3, 1); - zk.create(path, data.getBytes(), Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL, ver); - } else if (stat != null) { - zk.setData(path, data.getBytes(), -1); - } - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - } - - public List getChildren(String node, final boolean watch) { - - List childNodes = null; - - try { - childNodes = zk.getChildren(node, watch); - } catch (KeeperException | InterruptedException e) { - throw new IllegalStateException(e); - } - - return childNodes; - } - - public void watch4Data(String path) { - try { - Stat stat = new Stat(); - zk.getData(path, true, stat); - System.out.println(stat.getVersion()); - - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - - } - - public void watch4Children(String path) { - try { - Stat stat = new Stat(); - zk.getChildren(path, true); - - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - - } - - public void spinModel(String propFileAsString) { - // get childern for /nodes - // for each childern - get the data - // find which one to spin into - // get the lock - // spin - // release the lock - - try { - prop = new Properties(); - StringReader reader = new StringReader(propFileAsString); - - prop.load(reader); - - // get the name of the engine - String engineName = prop.getProperty(Constants.ENGINE); - String model = engineName; // prop.getProperty(Constants.MODEL); - - int numInstances = 1; - if (prop.containsKey(Settings.COUNT)) { - numInstances = Integer.parseInt(prop.getProperty(Settings.COUNT)); - } - - // the pattern is - // /model//engineName - // String nodeParent = "/model/" + model; - // List modelInstances = zk.getChildren(nodeParent, false); - - String modelPath = MODEL_ROOT + "/" + model; - - // check to see if the number of children is satisfied - // if(modelInstances.size() < numInstances) - { - // get all the stats - String needs = prop.getProperty(Settings.REQUIREMENTS); - Object needObject = gson.fromJson(needs, Object.class); // typically this is a dictionary - if (needObject instanceof Map) { - Map needMap = (Map) needObject; - Map capMap = getCurrentCapabilities(); - boolean accomodate = canAccomodate(capMap, needMap); - // take the mutex lock - getLock(modelPath); - - // force a fail - // accomodate = false; - if (accomodate) { - // check the status to see this is not something we got after someone else - // released it.. - String statusNode = MODEL_ROOT + "/" + model + "/status"; - String status = getNodeData(statusNode); - // make sure this is not released by someone else - if (status.equalsIgnoreCase("INIT")) { - // spin the server - // set the status - // connect2Py i.e. spin an engine - System.err.println("Spinning the server now.. "); - - // update the endpoint - updateNodeData(MODEL_ROOT + "/" + model + "/endpoint", "abracadabra", true); - - // update the status node - zk.setData(statusNode, "AVAILABLE".getBytes(), -1); - } - } else { - // report failure count - String failNode = MODEL_ROOT + "/" + model + "/fail/" + id; - // record this node - updateNodeData(failNode, "FAIL", true); - } - releaseLock(modelPath); - } - - } - } catch (NumberFormatException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (JsonSyntaxException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (IOException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - } - - public void getLock(String modelPath) { - // create the lock - // need a way to create node if one doesnt exist - try { - InterProcessMutex lock = new InterProcessMutex(client, modelPath + "/locker"); - boolean acquired = lock.acquire(3, TimeUnit.SECONDS); - if (acquired) { - modelLock.put(modelPath, lock); - } - } catch (Exception e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - } - - public void releaseLock(String modelPath) { - try { - if (modelLock.containsKey(modelPath)) { - InterProcessMutex lock = (InterProcessMutex) modelLock.get(modelPath); - lock.release(); - } - } catch (Exception e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - } - - public Map getCurrentCapabilities() { - boolean initialized = pyTranslator.getBoolean("'hardware_util' in globals()"); - if (!initialized) { - pyTranslator.runScript("import gaas_hardware_util as ghu"); - pyTranslator.runScript("hardware_util = ghu.HardwareUtil()"); - } - String capString = "" + pyTranslator.runScript("hardware_util.get_all()"); - System.err.println("Cap String .. " + capString); - Object cap = gson.fromJson(capString, Object.class); - if (cap instanceof Map) { - Map capMap = (Map) cap; - return capMap; - } else { - return new HashMap(); - } - - } - - public boolean canAccomodate(Map capMap, Map needMap) { - // compares the map with each of it - boolean accomodate = true; - Iterator keys = needMap.keySet().iterator(); - while (keys.hasNext() && accomodate) { - String thisNeed = keys.next(); - String needValueStr = needMap.get(thisNeed) + ""; - long needValue = Long.parseLong(needValueStr); - - // this could be a another map - Object itemMap = capMap.get(thisNeed); - if (itemMap instanceof Map) { - double capValue = Double.parseDouble(((Map) itemMap).get(AVAILABLE) + ""); - accomodate = accomodate && (needValue <= capValue); - } else { - accomodate = false; - } - } - return accomodate; - } - - public NativePySocketClient connect2Py(String modelPath) { - String port = PortAllocator.getInstance().getNextAvailablePort() + ""; - String timeout = "-1"; - - createCacheFolder(); - - // TODO verify this is correct. - String loggerLevel = this.prop.getProperty(Settings.LOGGER_LEVEL, "WARNING"); - Object[] outputs = ClientProcessWrapper.startTCPServerNativePy(this.workingDirectoryBasePath, port, null, - timeout, loggerLevel, true); - this.process = (Process) outputs[0]; - this.prefix = (String) outputs[1]; - - NativePySocketClient client = new NativePySocketClient(); - client.connect("127.0.0.1", Integer.parseInt(port), false); - - // connect the client - client = connectClient(client); - if (modelPath != null) { - modelClient.put(modelPath, client); - } - - return client; - } - - private void createCacheFolder() { - // create a generic folder - this.workingDirectory = "MODEL_" + Utility.getRandomString(6); - this.workingDirectoryBasePath = Utility.getInsightCacheDir() + "/" + workingDirectory; - this.cacheFolder = new File(workingDirectoryBasePath); - - // make the folder if one does not exist - if (!this.cacheFolder.exists()) { - this.cacheFolder.mkdir(); - } - } - - public NativePySocketClient connectClient(NativePySocketClient client) { - Thread t = new Thread(client); - t.start(); - while (!client.isReady()) { - synchronized (client) { - try { - client.wait(2000); - classLogger.info("Setting the socket client "); - } catch (InterruptedException e) { - classLogger.error(Constants.STACKTRACE, e); - } - } - } - return client; - } - - public void addCuratorListener(CuratorCacheListener listener, String path) { - CuratorCache tc = CuratorCache.build(client, path); - tc.start(); - tc.listenable().addListener(listener); - } - - public void catchup() { - // this is useful when the server starts up - try { - catchup = true; - List children = zk.getChildren(MODEL_ROOT, false); - for (int childIndex = 0; childIndex < children.size(); childIndex++) { - String thisModel = children.get(childIndex); - String modelPath = MODEL_ROOT + "/" + thisModel; - String statusNode = modelPath + "/status"; - - String status = getNodeData(statusNode); - if (status.equalsIgnoreCase("INIT")) { - String propFileAsString = getNodeData(modelPath); - spinModel(propFileAsString); - } - } - catchup = false; - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - - } - - @Override - public void event(Type type, ChildData oldData, ChildData data) { - // need to see if this is - - // TODO Auto-generated method stub - if (oldData != null) { - System.err.println("OLD >> " + oldData.getPath()); - } - if (data != null) { - System.err.println("NEW >> " + data.getPath()); - } - - // need to remove the base path to see if this is a parent level - // or some child level - // only react if it is a parent level - // i.e. /model/something.. - - if (data != null && data.getPath().endsWith("/status")) { - String inPath = data.getPath(); - String childPath = inPath.replace(MODEL_ROOT, ""); - String modelPath = data.getPath().replace("/status", ""); - { - try { - byte[] smssBytes = zk.getData(modelPath, false, new Stat()); - ModelInitListener scl = new ModelInitListener(modelPath, this); - - String propFileAsString = new String(smssBytes, "UTF-8"); - prop = new Properties(); - StringReader reader = new StringReader(propFileAsString); - prop.load(reader); - String modelId = prop.getProperty(Constants.ENGINE); - modelSMSS.put(MODEL_ROOT + "/" + modelId, propFileAsString); - - // dont spin the model - // wait for event on /status - zk.addWatch(modelPath, AddWatchMode.PERSISTENT_RECURSIVE); - String statusNode = MODEL_ROOT + "/" + modelId + "/status"; - String status = getNodeData(statusNode); - if (status != null && status.equalsIgnoreCase("INIT")) { - spinModel(propFileAsString); - // zk.addWatch(statusNode, this, AddWatchMode.PERSISTENT); - // spinModel(propFileAsString); - } - } catch (UnsupportedEncodingException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (KeeperException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (InterruptedException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } catch (IOException e) { - // TODO Auto-generated catch block - classLogger.error(Constants.STACKTRACE, e); - } - - } - } - } - - public void addZKListener(IModelZKListener listener) { - String path = listener.getPath(); - List listenerList = new Vector(); - if (listeners.containsKey(path)) { - listenerList = listeners.get(path); - } - listenerList.add(listener); - listeners.put(path, listenerList); - } - - public void processEvent(String path, WatchedEvent event) { - // get all the listeners for this path - // check to see if the event is listed - // if the event is listed.. check the predicate ? <-- not sure I need this given - // I am making all of it specific - if (listeners.containsKey(path)) { - List listenerList = listeners.get(path); - for (int listenerIndex = 0; listenerIndex < listenerList.size(); listenerIndex++) { - IModelZKListener thisListener = listenerList.get(listenerIndex); - - List events = thisListener.getEvents(); - if (events.contains(event.getType())) { - // process this event - thisListener.process(path, zk); - } - - } - } - } - -// public static void main(String [] args) -// { -// DIHelper helper = DIHelper.getInstance(); -// helper.loadCoreProp("C:/users/pkapaleeswaran/workspacej3/SemossDev/RDF_Map.prop"); -// ModelZKServer zkc = ModelZKServer.getInstance("RANDOM_ID"); -// -// -// InterProcessMutex lock = new InterProcessMutex(zkc.client, "/locker"); -// try { -// lock.acquire(); //(3, TimeUnit.SECONDS)) -// { -// try { -// // do some work inside of the critical section here -// System.err.println("Starting this thread "); -// Thread.sleep(10000); -// System.err.println("Going to release lock"); -// } finally { -// lock.release(); -// } -// } -// } catch (Exception e) { -// throw new RuntimeException(e); -// } -// try { -// Thread.sleep(10000000); -// } catch (InterruptedException e) { -// // TODO Auto-generated catch block -// classLogger.error(Constants.STACKTRACE, e); -// } -// -// //zkc.waitHere(); -// -// } - -} diff --git a/src/prerna/cluster/util/RemoteClientServerZK.java b/src/prerna/cluster/util/RemoteClientServerZK.java deleted file mode 100644 index 4ddacf5c8e5..00000000000 --- a/src/prerna/cluster/util/RemoteClientServerZK.java +++ /dev/null @@ -1,565 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; - -import org.apache.curator.framework.CuratorFramework; -import org.apache.curator.framework.CuratorFrameworkFactory; -import org.apache.curator.framework.recipes.cache.CuratorCache; -import org.apache.curator.framework.recipes.cache.CuratorCacheListener; -import org.apache.curator.framework.state.ConnectionState; -import org.apache.curator.retry.RetryOneTime; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONObject; - -import prerna.engine.api.RemoteModelStateEnum; - -/** - * This is a Singleton class used to exclusively manage a connection to - * ZooKeeper for tracking RemoteClientServer deployments on our cluster. This - * will track the state of each model deployment, including whether it is - * warming or active. This class handles holding requests until models are in an - * active state by performing health checks on the FastAPI service running in - * the container. This is required to give the FastAPI service a grace time - * between when the model is added to the active path and the time it requires - * to start up in the container. This is used by engines that extend the - * AbstractRemoteModelEngine IE: NEREngine.. - */ -public final class RemoteClientServerZK implements IRemoteClientServer { - - private static final Logger classLogger = LogManager.getLogger(RemoteClientServerZK.class); - - private static volatile RemoteClientServerZK instance; - - private static final String WARMING_PATH = "/models/warming"; - private static final String ACTIVE_PATH = "/models/active"; - private static final String MODEL_SCALER_PATH = "/services/kube-model-deployer"; - - // Connection-related fields - private CuratorFramework client; - private String zkServer = "localhost:2181"; - private boolean connected = false; - private Map env; - - // State tracking - private final ConcurrentMap modelStates = new ConcurrentHashMap<>(); - private final ConcurrentMap modelClusterIps = new ConcurrentHashMap<>(); - private final ConcurrentMap modelNames = new ConcurrentHashMap<>(); - private CuratorCache warmingCache; - private CuratorCache activeCache; - - public String modelScalerIp; - - private Boolean devPortForwarding = false; - - private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1); - - private RemoteClientServerZK() { - - } - - public static RemoteClientServerZK getInstance() { - if (instance != null) { - return instance; - } - - if (instance == null) { - synchronized (RemoteClientServerZK.class) { - if (instance == null) { - // Fully initialize before publishing to the volatile field so other - // threads never observe a half-constructed singleton (mid-init()). - RemoteClientServerZK proxy = new RemoteClientServerZK(); - proxy.init(); - instance = proxy; - classLogger.info("RemoteClientServerZK is initialized."); - } - } - } - return instance; - } - - private void init() { - try { - env = System.getenv(); - if (env.containsKey("ZK_SERVER")) { - zkServer = env.get("ZK_SERVER"); - } - - // Initialize Curator client - RetryOneTime retryPolicy = new RetryOneTime(1000); - client = CuratorFrameworkFactory.builder().connectString(zkServer).retryPolicy(retryPolicy) - .connectionTimeoutMs(5000).sessionTimeoutMs(10000).maxCloseWaitMs(2000).build(); - - // Add connection state listener - client.getConnectionStateListenable().addListener((client, state) -> { - classLogger.info("ZooKeeper connection state changed to: {}", state); - if (state == ConnectionState.CONNECTED || state == ConnectionState.RECONNECTED) { - connected = true; - } else if (state == ConnectionState.LOST || state == ConnectionState.SUSPENDED) { - connected = false; - } - }); - - client.start(); - - // Wait for connection - if (!client.blockUntilConnected(10, TimeUnit.SECONDS)) { - throw new IllegalStateException("Failed to connect to ZooKeeper"); - } - - setupZKPaths(); - setupCaches(); - - loadInitialState(); - - setupCacheRefresher(); - - } catch (Exception e) { - classLogger.error("Failed to initialize ZooKeeper connection", e); - throw new RuntimeException("Failed to initialize ZooKeeper connection", e); - } - } - - private void loadInitialState() { - try { - // Loading model scaler IP first - if (client.checkExists().forPath(MODEL_SCALER_PATH) != null) { - byte[] scalerData = client.getData().forPath(MODEL_SCALER_PATH); - if (scalerData != null && scalerData.length > 0) { - modelScalerIp = new String(scalerData, "UTF-8"); - classLogger.info("Discovered model scaler IP: {}", modelScalerIp); - } else { - classLogger.error("Model scaler path exists but no IP data found"); - } - } - - // Load active models - List activeModels = client.getChildren().forPath(ACTIVE_PATH); - for (String modelId : activeModels) { - String path = ACTIVE_PATH + "/" + modelId; - classLogger.info("Loading data for active model at path: {}", path); - - byte[] data = client.getData().forPath(path); - if (data != null && data.length > 0) { - String rawData = new String(data, "UTF-8"); - - JSONObject jsonData = new JSONObject(rawData); - String clusterIp = jsonData.getString("ip"); - String modelName = jsonData.getString("model_name"); - - modelStates.put(modelId, RemoteModelStateEnum.ACTIVE); - modelClusterIps.put(modelId, clusterIp); - modelNames.put(modelId, modelName); - - classLogger.info("Loaded active model {} ({}) with IP: {} from ZK.", modelId, modelName, clusterIp); - } - } - - // Load warming models - List warmingModels = client.getChildren().forPath(WARMING_PATH); - for (String modelId : warmingModels) { - modelStates.put(modelId, RemoteModelStateEnum.WARMING); - classLogger.info("Loaded warming model: {}", modelId); - } - } catch (Exception e) { - classLogger.error("Error loading initial state", e); - } - } - - private void setupCacheRefresher() { - classLogger.info("Setting up cache refresher for remote models..."); - scheduler.scheduleAtFixedRate(() -> { - try { - refreshModelStates(); - } catch (Exception e) { - classLogger.error("Error refreshing model states", e); - } - }, 30, 30, TimeUnit.SECONDS); // Refresh every 30 seconds - } - - private void refreshModelStates() throws Exception { - classLogger.debug("Refreshing model states..."); - - // Refresh active models - List activeModels = client.getChildren().forPath(ACTIVE_PATH); - for (String modelId : activeModels) { - modelStates.put(modelId, RemoteModelStateEnum.ACTIVE); - } - - // Refresh warming models - List warmingModels = client.getChildren().forPath(WARMING_PATH); - for (String modelId : warmingModels) { - if (!modelStates.get(modelId).equals(RemoteModelStateEnum.ACTIVE)) { - modelStates.put(modelId, RemoteModelStateEnum.WARMING); - } - } - - // Clean up stale entries - for (String modelId : new HashSet<>(modelStates.keySet())) { - if (!activeModels.contains(modelId) && !warmingModels.contains(modelId)) { - modelStates.put(modelId, RemoteModelStateEnum.COLD); - } - } - - classLogger.debug("Refreshed model states, current map: {}", modelStates); - } - - @Override - public String getModelScalerIp() { - return modelScalerIp; - } - - public void close() { - try { - if (warmingCache != null) { - warmingCache.close(); - } - if (activeCache != null) { - activeCache.close(); - } - if (client != null) { - client.close(); - } - } catch (Exception e) { - classLogger.error("Error closing resources", e); - } - } - - private void setupZKPaths() { - try { - // Ensure base paths exist - if (client.checkExists().forPath(WARMING_PATH) == null) { - client.create().creatingParentsIfNeeded().forPath(WARMING_PATH); - } - if (client.checkExists().forPath(ACTIVE_PATH) == null) { - client.create().creatingParentsIfNeeded().forPath(ACTIVE_PATH); - } - } catch (Exception e) { - classLogger.error("Error setting up ZK paths", e); - } - } - - private void setupCaches() { - try { - // Set up cache for warming path - warmingCache = CuratorCache.build(client, WARMING_PATH); - CuratorCacheListener warmingListener = CuratorCacheListener.builder().forCreates(node -> { - String modelId = getModelIdFromPath(node.getPath()); - classLogger.info("Model {} entered warming state", modelId); - modelStates.put(modelId, RemoteModelStateEnum.WARMING); - }).forDeletes(node -> { - String modelId = getModelIdFromPath(node.getPath()); - classLogger.info("Model {} left warming state", modelId); - // Only update if not active - if (!modelStates.get(modelId).equals(RemoteModelStateEnum.ACTIVE)) { - modelStates.put(modelId, RemoteModelStateEnum.COLD); - } - }).build(); - warmingCache.listenable().addListener(warmingListener); - warmingCache.start(); - - // Set up cache for active path - activeCache = CuratorCache.build(client, ACTIVE_PATH); - CuratorCacheListener activeListener = CuratorCacheListener.builder().forCreates(node -> { - String modelId = getModelIdFromPath(node.getPath()); - classLogger.info("Model {} became active, processing data from path: {}", modelId, node.getPath()); - modelStates.put(modelId, RemoteModelStateEnum.ACTIVE); - - try { - // First try immediate read - byte[] data = node.getData(); - if (data == null || data.length == 0) { - classLogger.debug("Initial data read empty for {}, starting retry sequence", modelId); - // If empty Retry - data = getNodeDataWithRetry(node.getPath(), 5, 500); // 5 retries, 500ms delay - } - - if (data == null || data.length == 0) { - classLogger.error("No data found for active model {} after retries", modelId); - return; - } - - String rawData = new String(data, "UTF-8"); - - classLogger.info("Processing raw data for model {}: {}", modelId, rawData); - - JSONObject jsonData = new JSONObject(rawData); - String clusterIp = jsonData.getString("ip"); - String modelName = jsonData.getString("model_name"); - - modelClusterIps.put(modelId, clusterIp); - modelNames.put(modelId, modelName); - - classLogger.info("Successfully registered model {} ({}) with IP: {}", modelId, modelName, - clusterIp); - } catch (Exception e) { - classLogger.error("Error processing data for model {}: {}", modelId, e.getMessage(), e); - } - }).forDeletes(node -> { - String modelId = getModelIdFromPath(node.getPath()); - String modelName = modelNames.get(modelId); - classLogger.info("Model {} ({}) is no longer active", modelId, modelName); - modelClusterIps.remove(modelId); - modelNames.remove(modelId); - if (isModelWarming(modelId)) { - modelStates.put(modelId, RemoteModelStateEnum.WARMING); - } else { - modelStates.put(modelId, RemoteModelStateEnum.COLD); - } - }).build(); - activeCache.listenable().addListener(activeListener); - activeCache.start(); - - } catch (Exception e) { - classLogger.error("Error setting up ZK caches", e); - } - } - - private byte[] getNodeDataWithRetry(String path, int maxRetries, long delayMs) { - for (int i = 0; i < maxRetries; i++) { - try { - byte[] data = client.getData().forPath(path); - if (data != null && data.length > 0) { - return data; - } - Thread.sleep(delayMs); - } catch (Exception e) { - classLogger.debug("Attempt {} to read data from {} failed: {}", i + 1, path, e.getMessage()); - } - } - return null; - } - - @Override - public String getModelName(String modelId) { - return modelNames.get(modelId); - } - - private String getModelIdFromPath(String path) { - return path.substring(path.lastIndexOf('/') + 1); - } - - @Override - public RemoteModelStateEnum getModelState(String modelId) { - return modelStates.getOrDefault(modelId, RemoteModelStateEnum.COLD); - } - - private String getModelPath(String modelId) { - return modelId; - } - - @Override - public String getModelClusterIp(String modelId) { - try { - String fullPath = ACTIVE_PATH + "/" + modelId; - if (client.checkExists().forPath(fullPath) == null) { - classLogger.error("Path does not exist: {}", fullPath); - return null; - } - - byte[] data = client.getData().forPath(fullPath); - - if (data == null || data.length == 0) { - classLogger.error("No data found at path: {}", fullPath); - return null; - } - - String rawData = new String(data, "UTF-8"); - - JSONObject jsonData = new JSONObject(rawData); - String ip = jsonData.getString("ip"); - - // Update cache maps - modelClusterIps.put(modelId, ip); - modelNames.put(modelId, jsonData.getString("model_name")); - - return ip; - } catch (Exception e) { - classLogger.error("Error getting cluster IP for model {}: {}", modelId, e.getMessage(), e); - return null; - } - } - - @Override - public boolean isModelWarming(String modelId) { - try { - String modelPath = getModelPath(modelId); - return client.checkExists().forPath(WARMING_PATH + "/" + modelPath) != null; - } catch (Exception e) { - classLogger.error("Error checking warming state for model {}", modelId, e); - return false; - } - } - - @Override - public boolean isModelActive(String modelId) { - try { - String modelPath = getModelPath(modelId); - return client.checkExists().forPath(ACTIVE_PATH + "/" + modelPath) != null; - } catch (Exception e) { - classLogger.error("Error checking active state for model {}", modelId, e); - return false; - } - } - - @Override - public boolean waitForModelActive(String modelId, long timeoutMs) { - try { - long startTime = System.currentTimeMillis(); - - // Waiting for model to appear in active path - while (System.currentTimeMillis() - startTime < timeoutMs) { - if (isModelActive(modelId)) { - classLogger.info("Model {} was found in the active path", modelId); - return true; - } else { - classLogger.info("Model {} in a warming wait loop..", modelId); - } - Thread.sleep(3000); - } - - classLogger.warn("Timeout waiting for model {} to appear in active path after {}ms", modelId, timeoutMs); - return false; - - } catch (Exception e) { - classLogger.error("Error waiting for model {} to become active", modelId, e); - return false; - } - } - - /** - * Waits for a model to reach a specific state, with a timeout - * - * @param modelId The ID of the model to wait for - * @param desiredState The state to wait for - * @param timeoutMs Maximum time to wait in milliseconds - * @return true if the model reached the desired state within the timeout, false - * otherwise - */ - @Override - public boolean waitForState(String modelId, RemoteModelStateEnum desiredState, long timeoutMs) { - try { - long startTime = System.currentTimeMillis(); - while (System.currentTimeMillis() - startTime < timeoutMs) { - RemoteModelStateEnum currentState = getModelState(modelId); - if (currentState == desiredState) { - return true; - } - // If we're waiting for ACTIVE but hit FAILED, break early - if (desiredState == RemoteModelStateEnum.ACTIVE && currentState == RemoteModelStateEnum.FAILED) { - classLogger.error("Model {} failed while waiting for active state", modelId); - return false; - } - Thread.sleep(1000); // 1 second - } - classLogger.warn("Timeout waiting for model {} to reach state {} after {}ms", modelId, desiredState, - timeoutMs); - return false; - } catch (Exception e) { - classLogger.error("Error waiting for model {} to reach state {}", modelId, desiredState, e); - return false; - } - } - - /** - * Gets a list of all active models with their associated information - * - * @return List of ModelInfo objects containing model details - */ - @Override - public List getActiveModels() { - List activeModels = new ArrayList<>(); - - try { - List activeModelIds = client.getChildren().forPath(ACTIVE_PATH); - - for (String modelId : activeModelIds) { - String name = modelNames.get(modelId); - RemoteModelStateEnum state = modelStates.getOrDefault(modelId, RemoteModelStateEnum.COLD); - - if (name != null) { - activeModels.add(new RemoteModelInfo(modelId, name, state)); - } else { - classLogger.warn("Incomplete data for active model {}: name={}", modelId, name); - - String fullPath = ACTIVE_PATH + "/" + modelId; - byte[] data = client.getData().forPath(fullPath); - if (data != null && data.length > 0) { - String rawData = new String(data, "UTF-8"); - JSONObject jsonData = new JSONObject(rawData); - - name = jsonData.getString("model_name"); - - activeModels.add(new RemoteModelInfo(modelId, name, state)); - - modelNames.put(modelId, name); - } - } - } - } catch (Exception e) { - classLogger.error("Error getting active models with details", e); - } - - return activeModels; - } - - /** - * Gets a list of all warming models with their associated information - * - * @return List of RemoteModelInfo objects containing model details - */ - @Override - public List getWarmingModels() { - List warmingModels = new ArrayList<>(); - - try { - List warmingModelIds = client.getChildren().forPath(WARMING_PATH); - - for (String modelId : warmingModelIds) { - String name = modelNames.get(modelId); - RemoteModelStateEnum state = modelStates.getOrDefault(modelId, RemoteModelStateEnum.WARMING); - - warmingModels.add(new RemoteModelInfo(modelId, name != null ? name : "Warming...", state)); - } - } catch (Exception e) { - classLogger.error("Error getting warming models with details", e); - } - - return warmingModels; - } -} \ No newline at end of file diff --git a/src/prerna/cluster/util/RemoteClientServerZKRESTProxy.java b/src/prerna/cluster/util/RemoteClientServerZKRESTProxy.java deleted file mode 100644 index 3437a7cfc96..00000000000 --- a/src/prerna/cluster/util/RemoteClientServerZKRESTProxy.java +++ /dev/null @@ -1,609 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; - -import org.apache.http.HttpEntity; -import org.apache.http.client.config.RequestConfig; -import org.apache.http.client.methods.CloseableHttpResponse; -import org.apache.http.client.methods.HttpGet; -import org.apache.http.client.methods.HttpHead; -import org.apache.http.impl.client.CloseableHttpClient; -import org.apache.http.impl.client.HttpClients; -import org.apache.http.util.EntityUtils; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONArray; -import org.json.JSONObject; - -import prerna.engine.api.RemoteModelStateEnum; - -/** - * This class provides the same functionality as RemoteClientServerZK but - * connects to the REST proxy for ZooKeeper instead of connecting directly to - * ZooKeeper. It is designed for local development environments when direct - * ZooKeeper access isn't available. - */ -public class RemoteClientServerZKRESTProxy implements IRemoteClientServer { - - private static final Logger classLogger = LogManager.getLogger(RemoteClientServerZKRESTProxy.class); - - private static volatile RemoteClientServerZKRESTProxy instance; - - private static final String WARMING_PATH = "/models/warming"; - private static final String ACTIVE_PATH = "/models/active"; - private static final String MODEL_SCALER_PATH = "/services/kube-model-deployer"; - - private String zkRestProxyBaseUrl; - - private final ConcurrentMap modelStates = new ConcurrentHashMap<>(); - private final ConcurrentMap modelClusterIps = new ConcurrentHashMap<>(); - private final ConcurrentMap modelNames = new ConcurrentHashMap<>(); - - // Periodic refreshes - private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1); - - private RemoteClientServerZKRESTProxy() { - - } - - public static RemoteClientServerZKRESTProxy getInstance() { - if (instance == null) { - synchronized (RemoteClientServerZKRESTProxy.class) { - if (instance == null) { - // Fully initialize before publishing to the volatile field so other - // threads never observe a half-constructed singleton (mid-init()). - RemoteClientServerZKRESTProxy proxy = new RemoteClientServerZKRESTProxy(); - proxy.init(); - instance = proxy; - classLogger.info("RemoteClientServerZKRestProxy is initialized."); - } - } - } - return instance; - } - - private void init() { - try { - Map env = System.getenv(); - if (env.containsKey("ZK_INGRESS")) { - zkRestProxyBaseUrl = env.get("ZK_INGRESS"); - if (!zkRestProxyBaseUrl.endsWith("/")) { - zkRestProxyBaseUrl += "/"; - } - classLogger.info("Using ZK_INGRESS from environment: {}", zkRestProxyBaseUrl); - } else { - throw new IllegalStateException("ZK_INGRESS environment variable is required for ZK REST Proxy mode"); - } - - // Validate connection to ZK REST Proxy - validateConnection(); - - // Initial load of model states - loadInitialState(); - - // Setup periodic refresh - setupCacheRefresher(); - - } catch (Exception e) { - classLogger.error("Failed to initialize ZooKeeper REST Proxy connection", e); - throw new RuntimeException("Failed to initialize ZooKeeper REST Proxy connection", e); - } - } - - private void validateConnection() throws Exception { - String healthCheckUrl = zkRestProxyBaseUrl + "health"; - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(5000).setSocketTimeout(5000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpGet httpGet = new HttpGet(healthCheckUrl); - - try (CloseableHttpResponse response = httpClient.execute(httpGet)) { - int statusCode = response.getStatusLine().getStatusCode(); - if (statusCode != 200) { - throw new IllegalStateException( - "ZooKeeper REST Proxy health check failed with status code: " + statusCode); - } - - classLogger.info("Successfully connected to ZooKeeper REST Proxy"); - } - } catch (IOException e) { - classLogger.error("Failed to connect to ZooKeeper REST Proxy", e); - throw new IllegalStateException("Failed to connect to ZooKeeper REST Proxy: " + e.getMessage(), e); - } - } - - /** - * Gets ZNode data as a JSONObject from the REST proxy response. This method - * extracts the nested "data" object from the response. - * - * @param path The ZNode path - * @return JSONObject containing the data field contents, or null if error/not - * found - */ - private JSONObject getZNodeDataAsJson(String path) { - String formattedPath = path; - if (formattedPath.startsWith("/")) { - formattedPath = formattedPath.substring(1); - } - - String url = zkRestProxyBaseUrl + "znodes/v1/" + formattedPath; - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(5000).setSocketTimeout(5000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpGet httpGet = new HttpGet(url); - - try (CloseableHttpResponse response = httpClient.execute(httpGet)) { - int statusCode = response.getStatusLine().getStatusCode(); - if (statusCode == 404) { - classLogger.warn("ZNode not found at path: {}", path); - return null; - } - - if (statusCode != 200) { - classLogger.error("Failed to get ZNode data, status code: {}", statusCode); - return null; - } - - HttpEntity entity = response.getEntity(); - if (entity != null) { - String jsonResponse = EntityUtils.toString(entity); - JSONObject responseObj = new JSONObject(jsonResponse); - - if (responseObj.has("data") && !responseObj.isNull("data")) { - // The data field is a nested JSON object - JSONObject dataObj = responseObj.getJSONObject("data"); - return dataObj; - } - } - } - - } catch (Exception e) { - classLogger.error("Error getting ZNode data for path {}: {}", path, e.getMessage()); - } - - return null; - } - - private void loadInitialState() { - try { - String modelScalerIpData = getZNodeData(MODEL_SCALER_PATH); - if (modelScalerIpData != null) { - classLogger.info("Model scaler IP from REST proxy: {}", modelScalerIpData); - } - - List activeModels = getZNodeChildren(ACTIVE_PATH); - for (String modelId : activeModels) { - String modelPath = ACTIVE_PATH + "/" + modelId; - classLogger.info("Loading data for active model at path: {}", modelPath); - - JSONObject modelData = getZNodeDataAsJson(modelPath); - if (modelData != null) { - String clusterIp = modelData.getString("ip"); - String modelName = modelData.getString("model_name"); - - modelStates.put(modelId, RemoteModelStateEnum.ACTIVE); - modelClusterIps.put(modelId, clusterIp); - modelNames.put(modelId, modelName); - - classLogger.info("Loaded active model {} ({}) with IP: {} from ZK REST Proxy.", modelId, modelName, - clusterIp); - } - } - - List warmingModels = getZNodeChildren(WARMING_PATH); - for (String modelId : warmingModels) { - modelStates.put(modelId, RemoteModelStateEnum.WARMING); - classLogger.info("Loaded warming model: {}", modelId); - } - - } catch (Exception e) { - classLogger.error("Error loading initial state from ZK REST Proxy", e); - } - } - - private void setupCacheRefresher() { - classLogger.info("Setting up cache refresher for remote models..."); - scheduler.scheduleAtFixedRate(() -> { - try { - refreshModelStates(); - } catch (Exception e) { - classLogger.error("Error refreshing model states", e); - } - }, 30, 30, TimeUnit.SECONDS); // Refresh every 30 seconds - } - - private void refreshModelStates() { - classLogger.debug("Refreshing model states from ZK REST Proxy..."); - - try { - List activeModels = getZNodeChildren(ACTIVE_PATH); - for (String modelId : activeModels) { - modelStates.put(modelId, RemoteModelStateEnum.ACTIVE); - } - - List warmingModels = getZNodeChildren(WARMING_PATH); - for (String modelId : warmingModels) { - if (!modelStates.getOrDefault(modelId, RemoteModelStateEnum.COLD).equals(RemoteModelStateEnum.ACTIVE)) { - modelStates.put(modelId, RemoteModelStateEnum.WARMING); - } - } - - for (String modelId : activeModels) { - String path = ACTIVE_PATH + "/" + modelId; - updateModelInfoFromZNode(modelId, path); - } - - // Clean up stale entries - for (String modelId : new ArrayList<>(modelStates.keySet())) { - if (!activeModels.contains(modelId) && !warmingModels.contains(modelId)) { - modelStates.put(modelId, RemoteModelStateEnum.COLD); - } - } - - classLogger.debug("Refreshed model states, current map: {}", modelStates); - - } catch (Exception e) { - classLogger.error("Error refreshing model states from ZK REST Proxy", e); - } - } - - private void updateModelInfoFromZNode(String modelId, String path) { - JSONObject modelData = getZNodeDataAsJson(path); - - if (modelData != null) { - try { - String clusterIp = modelData.getString("ip"); - String modelName = modelData.getString("model_name"); - - // Update cache maps - modelClusterIps.put(modelId, clusterIp); - modelNames.put(modelId, modelName); - - classLogger.debug("Updated model info for {}: name={}, ip={}", modelId, modelName, clusterIp); - } catch (Exception e) { - classLogger.error("Error extracting model data fields for {}: {}", modelId, e.getMessage()); - } - } - } - - private String getZNodeData(String path) { - String formattedPath = path; - if (formattedPath.startsWith("/")) { - formattedPath = formattedPath.substring(1); - } - - String url = zkRestProxyBaseUrl + "znodes/v1/" + formattedPath; - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(5000).setSocketTimeout(5000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpGet httpGet = new HttpGet(url); - - try (CloseableHttpResponse response = httpClient.execute(httpGet)) { - int statusCode = response.getStatusLine().getStatusCode(); - if (statusCode == 404) { - classLogger.warn("ZNode not found at path: {}", path); - return null; - } - - if (statusCode != 200) { - classLogger.error("Failed to get ZNode data, status code: {}", statusCode); - return null; - } - - HttpEntity entity = response.getEntity(); - if (entity != null) { - String jsonResponse = EntityUtils.toString(entity); - JSONObject responseObj = new JSONObject(jsonResponse); - - if (responseObj.has("data") && !responseObj.isNull("data")) { - return responseObj.getString("data"); - } - } - } - - } catch (Exception e) { - classLogger.error("Error getting ZNode data for path {}: {}", path, e.getMessage()); - } - - return null; - } - - private List getZNodeChildren(String path) { - List children = new ArrayList<>(); - - String formattedPath = path; - if (formattedPath.startsWith("/")) { - formattedPath = formattedPath.substring(1); - } - - String url = zkRestProxyBaseUrl + "znodes/v1/" + formattedPath + "?view=children"; - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(5000).setSocketTimeout(5000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpGet httpGet = new HttpGet(url); - - try (CloseableHttpResponse response = httpClient.execute(httpGet)) { - int statusCode = response.getStatusLine().getStatusCode(); - if (statusCode == 404) { - classLogger.warn("ZNode not found at path: {}", path); - return children; - } - - if (statusCode != 200) { - classLogger.error("Failed to get ZNode children, status code: {}", statusCode); - return children; - } - - HttpEntity entity = response.getEntity(); - if (entity != null) { - String jsonResponse = EntityUtils.toString(entity); - JSONObject responseObj = new JSONObject(jsonResponse); - - if (responseObj.has("children")) { - JSONArray childrenArray = responseObj.getJSONArray("children"); - for (int i = 0; i < childrenArray.length(); i++) { - children.add(childrenArray.getString(i)); - } - } - } - } - - } catch (Exception e) { - classLogger.error("Error getting ZNode children for path {}: {}", path, e.getMessage()); - } - - return children; - } - - private boolean checkZNodeExists(String path) { - String formattedPath = path; - if (formattedPath.startsWith("/")) { - formattedPath = formattedPath.substring(1); - } - - String url = zkRestProxyBaseUrl + "znodes/v1/" + formattedPath; - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(3000).setSocketTimeout(3000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpHead httpHead = new HttpHead(url); - - try (CloseableHttpResponse response = httpClient.execute(httpHead)) { - return response.getStatusLine().getStatusCode() == 200; - } - - } catch (Exception e) { - classLogger.error("Error checking if ZNode exists at path {}: {}", path, e.getMessage()); - } - - return false; - } - - public void close() { - try { - if (scheduler != null && !scheduler.isShutdown()) { - scheduler.shutdown(); - try { - if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) { - scheduler.shutdownNow(); - } - } catch (InterruptedException e) { - scheduler.shutdownNow(); - Thread.currentThread().interrupt(); - } - } - - } catch (Exception e) { - classLogger.error("Error closing resources", e); - } - } - - @Override - public RemoteModelStateEnum getModelState(String modelId) { - return modelStates.getOrDefault(modelId, RemoteModelStateEnum.COLD); - } - - @Override - public String getModelClusterIp(String modelId) { - String clusterIp = modelClusterIps.get(modelId); - - if (clusterIp == null) { - String path = ACTIVE_PATH + "/" + modelId; - JSONObject modelData = getZNodeDataAsJson(path); - - if (modelData != null && !modelData.isEmpty()) { - try { - classLogger.info("modelData: {}", modelData); - clusterIp = modelData.getString("ip"); - String modelName = modelData.getString("model_name"); - - modelClusterIps.put(modelId, clusterIp); - modelNames.put(modelId, modelName); - } catch (Exception e) { - classLogger.error("Error parsing model data for {}: {}", modelId, e.getMessage()); - } - } - } - - return clusterIp; - } - - @Override - public String getModelName(String modelId) { - return modelNames.get(modelId); - } - - @Override - public boolean isModelWarming(String modelId) { - String path = WARMING_PATH + "/" + modelId; - return checkZNodeExists(path); - } - - @Override - public boolean isModelActive(String modelId) { - String path = ACTIVE_PATH + "/" + modelId; - return checkZNodeExists(path); - } - - @Override - public boolean waitForModelActive(String modelId, long timeoutMs) { - try { - long startTime = System.currentTimeMillis(); - - while (System.currentTimeMillis() - startTime < timeoutMs) { - if (isModelActive(modelId)) { - classLogger.info("Model {} was found in the active path", modelId); - modelStates.put(modelId, RemoteModelStateEnum.ACTIVE); - return true; - } else { - classLogger.info("Model {} in a warming wait loop..", modelId); - } - Thread.sleep(3000); - } - classLogger.warn("Timeout waiting for model {} to appear in active path after {}ms", modelId, timeoutMs); - return false; - - } catch (Exception e) { - classLogger.error("Error waiting for model {} to become active", modelId, e); - return false; - } - } - - @Override - public boolean waitForState(String modelId, RemoteModelStateEnum desiredState, long timeoutMs) { - try { - long startTime = System.currentTimeMillis(); - while (System.currentTimeMillis() - startTime < timeoutMs) { - RemoteModelStateEnum currentState = getModelState(modelId); - if (currentState == desiredState) { - return true; - } - - // If we're waiting for ACTIVE but the model is COLD, try to check active path - // directly - if (desiredState == RemoteModelStateEnum.ACTIVE && currentState == RemoteModelStateEnum.COLD) { - if (isModelActive(modelId)) { - modelStates.put(modelId, RemoteModelStateEnum.ACTIVE); - return true; - } - } - - Thread.sleep(1000); // 1 second - } - - classLogger.warn("Timeout waiting for model {} to reach state {} after {}ms", modelId, desiredState, - timeoutMs); - return false; - - } catch (Exception e) { - classLogger.error("Error waiting for model {} to reach state {}", modelId, desiredState, e); - return false; - } - } - - @Override - public List getActiveModels() { - List activeModels = new ArrayList<>(); - - try { - List activeModelIds = getZNodeChildren(ACTIVE_PATH); - - for (String modelId : activeModelIds) { - String name = modelNames.get(modelId); - RemoteModelStateEnum state = modelStates.getOrDefault(modelId, RemoteModelStateEnum.ACTIVE); - - if (name != null) { - activeModels.add(new RemoteModelInfo(modelId, name, state)); - } else { - String path = ACTIVE_PATH + "/" + modelId; - JSONObject modelData = getZNodeDataAsJson(path); - - if (modelData != null) { - name = modelData.getString("model_name"); - - activeModels.add(new RemoteModelInfo(modelId, name, state)); - - modelNames.put(modelId, name); - } else { - activeModels.add(new RemoteModelInfo(modelId, "Unknown", state)); - } - } - } - - } catch (Exception e) { - classLogger.error("Error getting active models with details", e); - } - - return activeModels; - } - - @Override - public List getWarmingModels() { - List warmingModels = new ArrayList<>(); - - try { - List warmingModelIds = getZNodeChildren(WARMING_PATH); - - for (String modelId : warmingModelIds) { - String name = modelNames.get(modelId); - RemoteModelStateEnum state = modelStates.getOrDefault(modelId, RemoteModelStateEnum.WARMING); - - warmingModels.add(new RemoteModelInfo(modelId, name != null ? name : "Warming...", state)); - } - - } catch (Exception e) { - classLogger.error("Error getting warming models with details", e); - } - - return warmingModels; - } - - @Override - public String getModelScalerIp() { - return System.getenv("KMS_INGRESS"); - } -} \ No newline at end of file diff --git a/src/prerna/cluster/util/RemoteModelInfo.java b/src/prerna/cluster/util/RemoteModelInfo.java deleted file mode 100644 index 2030b278436..00000000000 --- a/src/prerna/cluster/util/RemoteModelInfo.java +++ /dev/null @@ -1,50 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import prerna.engine.api.RemoteModelStateEnum; - -/** - * Represents information about a model deployed in the remote infrastructure. - */ -public class RemoteModelInfo { - private final String id; - private final String name; - private final RemoteModelStateEnum state; - - public RemoteModelInfo(String id, String name, RemoteModelStateEnum state) { - this.id = id; - this.name = name; - this.state = state; - } - - // Getters - public String getId() { return id; } - public String getName() { return name; } - public RemoteModelStateEnum getState() { return state; } -} \ No newline at end of file diff --git a/src/prerna/cluster/util/ZKClientFactory.java b/src/prerna/cluster/util/ZKClientFactory.java deleted file mode 100644 index f98fa4edca4..00000000000 --- a/src/prerna/cluster/util/ZKClientFactory.java +++ /dev/null @@ -1,64 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -/** - * Factory class to get the appropriate ZooKeeper client implementation - * based on the environment configuration. - */ -public class ZKClientFactory { - - private static final Logger classLogger = LogManager.getLogger(ZKClientFactory.class); - - /** - * Gets the appropriate ZooKeeper client implementation based on environment variables. - * - * @return An implementation of IRemoteClientServer - */ - public static IRemoteClientServer getZKClient(Boolean devPortForwarding) { - String zkIngress = System.getenv("ZK_INGRESS"); - - if (zkIngress != null && !zkIngress.isEmpty() && !devPortForwarding) { - classLogger.info("ZK_INGRESS found, using REST proxy connection for ZooKeeper"); - - // When using ZK REST proxy, KMS_INGRESS is required - String kmsIngress = System.getenv("KMS_INGRESS"); - if (kmsIngress == null || kmsIngress.isEmpty()) { - classLogger.warn("KMS_INGRESS environment variable is not set but required for ZK REST Proxy mode"); - } - - return RemoteClientServerZKRESTProxy.getInstance(); - } else { - classLogger.info("Using direct ZooKeeper connection (cluster mode)"); - return RemoteClientServerZK.getInstance(); - } - } -} \ No newline at end of file diff --git a/src/prerna/engine/api/ModelTypeEnum.java b/src/prerna/engine/api/ModelTypeEnum.java index 5c40cbd8368..a3088cb1510 100644 --- a/src/prerna/engine/api/ModelTypeEnum.java +++ b/src/prerna/engine/api/ModelTypeEnum.java @@ -31,12 +31,7 @@ import prerna.engine.impl.model.AzureOpenAiEngine; import prerna.engine.impl.model.BedrockEngine; import prerna.engine.impl.model.EmbeddedModelEngine; -import prerna.engine.impl.model.KServeImageEmbedEngine; -import prerna.engine.impl.model.KServeImageEngine; -import prerna.engine.impl.model.KServeTTSEngine; -import prerna.engine.impl.model.KServeVisionEngine; import prerna.engine.impl.model.ModelRouterEngine; -import prerna.engine.impl.model.NEREngine; import prerna.engine.impl.model.OpenAiEngine; import prerna.engine.impl.model.TextEmbeddingsEngine; import prerna.engine.impl.model.TextGenerationEngine; @@ -58,11 +53,6 @@ public enum ModelTypeEnum { // these are secondary EMBEDDED("EMBEDDED", EmbeddedModelEngine.class.getName()), - NER("NER", NEREngine.class.getName()), - KSERVE_VISION("KSERVE_VISION", KServeVisionEngine.class.getName()), - KSERVE_IMAGE_EMBED("KSERVE_IMAGE_EMBED", KServeImageEmbedEngine.class.getName()), - KSERVE_IMAGE("KSERVE_IMAGE", KServeImageEngine.class.getName()), - KSERVE_TTS("KSERVE_TTS", KServeTTSEngine.class.getName()), REMOTE("REMOTE", RemoteModelEngine.class.getName()), TEXT_EMBEDDINGS("TEXT_EMBEDDINGS", TextEmbeddingsEngine.class.getName()), TEXT_GENERATION("TEXT_GENERATION", TextGenerationEngine.class.getName()), diff --git a/src/prerna/engine/api/RemoteModelStateEnum.java b/src/prerna/engine/api/RemoteModelStateEnum.java deleted file mode 100644 index 7587677d435..00000000000 --- a/src/prerna/engine/api/RemoteModelStateEnum.java +++ /dev/null @@ -1,36 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.api; - -public enum RemoteModelStateEnum { - COLD, - WARMING, - ACTIVE, - FAILED, - UNKNOWN -} diff --git a/src/prerna/engine/impl/function/OpenAITranscribeFunctionEngine.java b/src/prerna/engine/impl/function/OpenAITranscribeFunctionEngine.java index a6f1ffd2cef..e606a14b79d 100644 --- a/src/prerna/engine/impl/function/OpenAITranscribeFunctionEngine.java +++ b/src/prerna/engine/impl/function/OpenAITranscribeFunctionEngine.java @@ -41,17 +41,12 @@ import java.util.Properties; import java.util.UUID; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; - import prerna.engine.api.FunctionTypeEnum; import prerna.om.Insight; import prerna.util.Utility; public class OpenAITranscribeFunctionEngine extends AbstractFunctionEngine { - private static final ObjectMapper MAPPER = new ObjectMapper(); - public static final String URL = "URL"; public static final String API_KEY = "API_KEY"; public static final String MODEL = "MODEL"; @@ -59,7 +54,6 @@ public class OpenAITranscribeFunctionEngine extends AbstractFunctionEngine { private String url; private String apiKey; private String model; - private String modelType; @Override public void open(Properties smssProp) throws Exception { @@ -68,7 +62,6 @@ public void open(Properties smssProp) throws Exception { this.apiKey = smssProp.getProperty(API_KEY); this.url = smssProp.getProperty(URL); this.model = smssProp.getProperty(MODEL); - this.modelType = smssProp.getProperty("MODEL_TYPE"); if (this.apiKey == null || (this.apiKey.isEmpty())) { throw new RuntimeException("Must set API key. Use EMPTY if none."); @@ -80,11 +73,6 @@ public void open(Properties smssProp) throws Exception { @Override public Object execute(Map parameterValues) { - Boolean healthCheck = performHealthCheck(); - if (!healthCheck) { - throw new RuntimeException("Model " + this.model + "is offline."); - } - Insight insight = (Insight) parameterValues.getOrDefault("INSIGHT", null); String instanceDir = Utility.normalizePath(insight.getInsightFolder()); @@ -192,48 +180,6 @@ private static String extractTextField(String json) { return json; } - private static String buildHealthUrl(String rawUrl) { - if (rawUrl == null || rawUrl.isEmpty()) { - throw new IllegalArgumentException("URL is required for health check"); - } - - String base = rawUrl.endsWith("/") ? rawUrl.substring(0, rawUrl.length() - 1) : rawUrl; - - base = base.replaceFirst("/v1/?$", ""); - - return base + "/v2/health/live"; - } - - private Boolean performHealthCheck() { - if (this.modelType == null || !this.modelType.equalsIgnoreCase("kserve")) { - return true; - } - - try { - String healthUrl = buildHealthUrl(this.url); - - HttpClient client = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build(); - - HttpRequest request = HttpRequest.newBuilder().uri(URI.create(healthUrl)).timeout(Duration.ofSeconds(10)) - .GET().build(); - - HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); - - if (response.statusCode() / 100 != 2) { - throw new RuntimeException("Health check failed: HTTP " + response.statusCode()); - } - - JsonNode node = MAPPER.readTree(response.body()); - if (node.has("live") && node.get("live").asBoolean()) { - return true; - } else { - throw new RuntimeException("Model not online at " + healthUrl + ". Response: " + response.body()); - } - } catch (Exception e) { - throw new RuntimeException("Health check failed: " + e.getMessage(), e); - } - } - @Override public void close() throws IOException { // TODO Auto-generated method stub diff --git a/src/prerna/engine/impl/model/AbstractRemoteModelEngine.java b/src/prerna/engine/impl/model/AbstractRemoteModelEngine.java deleted file mode 100644 index 0606e6660b5..00000000000 --- a/src/prerna/engine/impl/model/AbstractRemoteModelEngine.java +++ /dev/null @@ -1,539 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model; - -import java.io.BufferedReader; -import java.io.IOException; -import java.io.InputStreamReader; -import java.nio.charset.StandardCharsets; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import java.util.concurrent.CompletableFuture; - -import org.apache.http.HttpEntity; -import org.apache.http.client.config.RequestConfig; -import org.apache.http.client.methods.CloseableHttpResponse; -import org.apache.http.client.methods.HttpPost; -import org.apache.http.entity.ContentType; -import org.apache.http.entity.StringEntity; -import org.apache.http.impl.client.CloseableHttpClient; -import org.apache.http.impl.client.HttpClients; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONObject; - -import prerna.cluster.util.IRemoteClientServer; -import prerna.cluster.util.RemoteClientServerZK; -import prerna.cluster.util.RemoteClientServerZKRESTProxy; -import prerna.cluster.util.ZKClientFactory; -import prerna.engine.api.ModelTypeEnum; -import prerna.engine.api.RemoteModelStateEnum; -import prerna.engine.impl.model.kserve.KServeAdapter; -import prerna.engine.impl.model.message.InputMessage; -import prerna.engine.impl.model.responses.AskModelEngineResponse; -import prerna.engine.impl.model.responses.EmbeddingsModelEngineResponse; -import prerna.om.Insight; -import prerna.util.Constants; -import prerna.util.Settings; - -/** - * This is a class used to be extended by models running on KServe. It contains - * methods for deploying the model to the cluster and making HTTP requests to - * the model. See https://github.com/SEMOSS/kubernetes-model-scaler for - * Kubernetes model scaling. - */ - -public class AbstractRemoteModelEngine extends AbstractModelEngine { - - private static final Logger classLogger = LogManager.getLogger(AbstractRemoteModelEngine.class); - - protected String model; - protected String modelRepoId; - protected String modelType; - private IRemoteClientServer zkClient; - // Use this to simulate the cluster environment - private Boolean devPortForwarding = false; - // For normal development - private String kmsIngressUrl = null; - private String modelIngressUrl = null; - private AbstractModelEngine implementingEngineClass = null; - - private final String INIT_PREFIX = "INIT_"; - - private enum Services { - KMS_START, // Kubernetes Model Scaler Start - KMS_SHUTDOWN, // Kubernetes Model Scaler Shutdown - MODEL // Model Specific Deployment - } - - @Override - public void open(Properties smssProp) throws Exception { - super.open(smssProp); - - if (this.smssProp.containsKey(Settings.MODEL)) { - this.model = this.smssProp.getProperty(Settings.MODEL).trim(); - } else { - throw new IllegalArgumentException("Model is not defined in SMSS file."); - } - - if (this.smssProp.containsKey(Settings.MODEL_REPO_ID)) { - this.modelRepoId = this.smssProp.getProperty(Settings.MODEL_REPO_ID).trim(); - } else { - throw new IllegalArgumentException("Model Repo ID is not defined in SMSS file."); - } - - if (this.smssProp.containsKey(Settings.MODEL_TYPE)) { - this.modelType = this.smssProp.getProperty(Settings.MODEL_TYPE).trim(); - } else { - throw new IllegalArgumentException("Model Type is not defined in SMSS file."); - } - - // Get the appropriate ZK client implementation based on environment - this.zkClient = ZKClientFactory.getZKClient(this.devPortForwarding); - - // Check if we're using the REST proxy (for KMS_INGRESS validation) - boolean usingRestProxy = this.zkClient instanceof RemoteClientServerZKRESTProxy; - - this.kmsIngressUrl = System.getenv("KMS_INGRESS"); - if (this.kmsIngressUrl != null && !this.kmsIngressUrl.isEmpty()) { - classLogger.info("Using KMS_INGRESS from environment: {}", this.kmsIngressUrl); - if (!this.kmsIngressUrl.endsWith("/")) { - this.kmsIngressUrl += "/"; - } - } else if (this.devPortForwarding) { - classLogger.info("Using devPortforwarding for KMS URL with localhost:8000/"); - } else { - classLogger.info( - "KMS_INGRESS environment variable not found and devPortforwarding not set, using ZooKeeper for KMS IP resolution. This is correct for production deployments."); - } - - this.modelIngressUrl = System.getenv("MODEL_INGRESS"); - if (this.modelIngressUrl != null && !this.modelIngressUrl.isEmpty()) { - classLogger.info("Using MODEL_INGRESS from environment: {}", this.modelIngressUrl); - if (!this.modelIngressUrl.endsWith("/")) { - this.modelIngressUrl += "/"; - } - } else if (this.devPortForwarding) { - classLogger.info("Using devPortForwarding for model URLs with localhost:8888/"); - } else { - classLogger.info( - "MODEL_INGRESS environment variable not found and devPortforwarding not set, using ZooKeeper for Model IP resolution. This is correct for production deployments."); - } - - String initEngineTypeKey = INIT_PREFIX + Constants.ENGINE_TYPE; - String initEngineType = smssProp.getProperty(initEngineTypeKey); - - if (initEngineType != null && !initEngineType.isEmpty()) { - implementingEngineClass = (AbstractModelEngine) Class.forName(initEngineType).newInstance(); - Properties implEngineSmss = new Properties(); - for (Object key : smssProp.keySet()) { - String keyStr = (String) key; - if (keyStr.equals(Constants.ENGINE_TYPE)) { - implEngineSmss.put(Constants.ENGINE_TYPE, initEngineType); - } else { - implEngineSmss.put(keyStr, smssProp.getProperty(keyStr)); - } - } - String engineId = smssProp.getProperty(Constants.ENGINE); - implementingEngineClass.open(implEngineSmss); - } - } - - private String createServiceUrl(Services service) throws Exception { - // Priority order: - // 1. Use devPortForwarding if enabled - // 2. Use kmsIngressUrl if available - // 3. Use service ip from ZooKeeper - String serviceUrl; - - // KMS SHUTDOWN - if (service == Services.KMS_SHUTDOWN) { - if (devPortForwarding) { - serviceUrl = String.format("http://localhost:8000/api/v2/stop?model_id=%s&model=%s", this.engineId, - this.model); - } else if (kmsIngressUrl != null) { - serviceUrl = String.format("%sapi/v2/stop?model_id=%s&model=%s", kmsIngressUrl, this.engineId, - this.model); - } else { - if (zkClient instanceof RemoteClientServerZKRESTProxy) { - throw new IllegalStateException( - "KMS_INGRESS environment variable must be set when using ZK REST Proxy"); - } - - RemoteClientServerZK directZkClient = (RemoteClientServerZK) zkClient; - String modelScalerIp = directZkClient.getModelScalerIp(); - if (modelScalerIp == null) { - classLogger.error("Unable to get model scaler IP from ZooKeeper"); - throw new IllegalStateException( - "Unable to get model scaler IP from ZooKeeper for shutdown operation."); - } - serviceUrl = String.format("http://%s/api/v2/stop?model_id=%s&model=%s", modelScalerIp, this.engineId, - this.model); - } - // KMS START - } else if (service == Services.KMS_START) { - if (devPortForwarding) { - serviceUrl = "http://localhost:8000/api/v2/start"; - } else if (kmsIngressUrl != null) { - serviceUrl = kmsIngressUrl + "api/v2/start"; - } else { - if (zkClient instanceof RemoteClientServerZKRESTProxy) { - throw new IllegalStateException( - "KMS_INGRESS environment variable must be set when using ZK REST Proxy"); - } - - RemoteClientServerZK directZkClient = (RemoteClientServerZK) zkClient; - String modelScalerIp = directZkClient.getModelScalerIp(); - if (modelScalerIp == null) { - classLogger.error("Unable to get model scaler IP from ZooKeeper"); - throw new IllegalStateException( - "Unable to get model scaler IP from ZooKeeper for deployment operation."); - } - serviceUrl = String.format("http://%s/api/v2/start", modelScalerIp); - } - // MODEL INFERENCE - } else if (service == Services.MODEL) { - // Grabbing the cluster IP in all situations since it should always have one - String clusterIp = zkClient.getModelClusterIp(this.engineId); - Boolean isModelTypeOpenAI = this.modelType.equals("OPEN_AI"); - if (clusterIp == null) { - classLogger.error("No cluster IP available for model {}", this.engineId); - throw new IllegalStateException("Unable to get cluster ip for model."); - } - // LOCAL DEV W/ PF - if (devPortForwarding) { - if (isModelTypeOpenAI) { - serviceUrl = "http://localhost:8080/openai/v1"; - } else { - serviceUrl = String.format("http://localhost:8080/v2/models/%s/infer", this.model); - } - // LOCAL DEV W/ INGRESS - } else if (this.modelIngressUrl != null) { - if (isModelTypeOpenAI) { - serviceUrl = this.modelIngressUrl + this.model + "/openai/v1"; - } else { - serviceUrl = this.modelIngressUrl + this.model + "/v2/models/" + this.model + "/infer"; - } - } - // DEPLOYMENT - else { - if (isModelTypeOpenAI) { - serviceUrl = String.format("http://%s/openai/v1", clusterIp); - } else { - serviceUrl = String.format("http://%s/v2/models/%s/infer", clusterIp, this.model); - } - } - } else { - throw new IllegalArgumentException("Unsupported service: " + service); - } - - return serviceUrl; - } - - // STARTING A MODEL.. - public boolean initiateAndWaitForDeployment(long timeoutMs) throws Exception { - if (zkClient.isModelActive(this.engineId)) { - classLogger.info("Model {} is already active", this.engineId); - return true; - } - - if (zkClient.isModelWarming(this.engineId)) { - classLogger.info("Model {} is already warming, waiting for activation", this.engineId); - return zkClient.waitForModelActive(this.engineId, timeoutMs); - } - - String deploymentUrl = this.createServiceUrl(Services.KMS_START); - classLogger.info("Using deployment URL: {}", deploymentUrl); - - // Deployment request in separate thread - CompletableFuture deploymentFuture = CompletableFuture.runAsync(() -> { - try (CloseableHttpClient httpClient = HttpClients.createDefault()) { - HttpPost httpPost = new HttpPost(deploymentUrl); - - JSONObject payload = new JSONObject(); - payload.put("model_id", this.engineId); - payload.put("model", this.model); - payload.put("model_repo_id", this.modelRepoId); - payload.put("model_type", this.modelType); - - StringEntity entity = new StringEntity(payload.toString(), ContentType.APPLICATION_JSON); - httpPost.setEntity(entity); - - httpClient.execute(httpPost).close(); - } catch (Exception e) { - // I'm not hanging the main thread on this request, I'll just monitor ZK for the - // model status - classLogger.warn( - "HTTP request to model scaler stilling progress, dropping connection but continuing to check ZooKeeper status", - e); - } - }); - - long startTime = System.currentTimeMillis(); - long warmingTimeout = Math.min(120000, timeoutMs); - - while (System.currentTimeMillis() - startTime < warmingTimeout) { - if (zkClient.isModelWarming(this.engineId)) { - classLogger.info("Model {} has entered warming state, waiting for activation", this.engineId); - // Canceling the HTTP request if still running - deploymentFuture.cancel(true); - return zkClient.waitForModelActive(this.engineId, timeoutMs - (System.currentTimeMillis() - startTime)); - } - Thread.sleep(1000); // 1 sec polling - } - - classLogger.error("Timeout waiting for model {} to enter warming state", this.engineId); - return false; - } - - public String shutdownModelRequest() throws Exception { - RemoteModelStateEnum currentState = zkClient.getModelState(this.engineId); - - if (currentState == RemoteModelStateEnum.COLD) { - classLogger.info("Model {} is already shutdown", this.engineId); - return String.format("Model %s is already shutdown", this.engineId); - } - - String shutdownUrl = this.createServiceUrl(Services.KMS_SHUTDOWN); - - classLogger.debug("Using KMS shutdown URL: {}", shutdownUrl); - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(30000).setSocketTimeout(30000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpPost httpPost = new HttpPost(shutdownUrl); - - try (CloseableHttpResponse response = httpClient.execute(httpPost)) { - int statusCode = response.getStatusLine().getStatusCode(); - if (statusCode != 200) { - String error = String.format("Shutdown request failed with status code: %d", statusCode); - classLogger.error(error); - return error; - } - - String message = String.format("Successfully initiated shutdown for model %s", this.engineId); - classLogger.info(message); - return message; - } - } catch (Exception e) { - String error = String.format("Error making shutdown request for model %s: %s", this.engineId, - e.getMessage()); - classLogger.error(error, e); - return error; - } - } - - protected JSONObject makeModelRequest(JSONObject requestPayload) throws Exception { - // Get current state and handle warming/cold states - RemoteModelStateEnum currentState = zkClient.getModelState(this.engineId); - - // If cold try deploy - if (currentState == RemoteModelStateEnum.COLD) { - boolean deployed = initiateAndWaitForDeployment(120000); // 2 min - if (!deployed) { - classLogger.error("Failed to deploy model {}", this.engineId); - return null; - } - currentState = zkClient.getModelState(this.engineId); - } - - // Always wait for active state whether it started as WARMING or just became - // WARMING after deployment - if (currentState == RemoteModelStateEnum.WARMING) { - classLogger.info("Model {} is warming, waiting for activation...", this.engineId); - boolean becameActive = zkClient.waitForModelActive(this.engineId, 300000); // 5 min - if (!becameActive) { - classLogger.error("Model {} failed to become active after warming", this.engineId); - return null; - } - currentState = zkClient.getModelState(this.engineId); - } - - if (currentState != RemoteModelStateEnum.ACTIVE) { - // TEMP SOLUTION TO CURATOR DESYNC ISSUE - classLogger.error( - "Model {} is not active in conncurrent hashmap. Current state: {}. Checking path directly", - this.engineId, currentState); - Boolean modelActive = zkClient.isModelActive(this.engineId); - if (!modelActive) { - classLogger.error("Model {} is not active in ZooKeeper", this.engineId); - return null; - } - } - - String clusterIp = zkClient.getModelClusterIp(this.engineId); - if (clusterIp == null) { - classLogger.error("No cluster IP available for model {}", this.engineId); - return null; - } - - return makeInferenceRequest(clusterIp, requestPayload); - } - - // For models that don't go through the OpenAI API (ie. NER, etc.) - private JSONObject makeInferenceRequest(String clusterIp, JSONObject requestPayload) throws Exception { - // Formatting the payload into KServe Protocol format - JSONObject kservePayload = KServeAdapter.toKServeRequest(requestPayload); - - classLogger.debug("Sending KServe payload: {}", kservePayload.toString(2)); - classLogger.info("Sending request to model {} at cluster IP {}", this.engineId, clusterIp); - - String url = this.createServiceUrl(Services.MODEL); - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(30000).setSocketTimeout(900000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpPost httpPost = new HttpPost(url); - httpPost.setHeader("Accept", "text/event-stream"); - httpPost.setHeader("Content-Type", "application/json"); - httpPost.setHeader("Inference-Header-Content-Length", "2000"); - - StringEntity entity = new StringEntity(kservePayload.toString(), ContentType.APPLICATION_JSON); - httpPost.setEntity(entity); - - try (CloseableHttpResponse response = httpClient.execute(httpPost)) { - int statusCode = response.getStatusLine().getStatusCode(); - if (statusCode != 200) { - classLogger.error("Request failed with status code: {}", statusCode); - return null; - } - - HttpEntity responseEntity = response.getEntity(); - if (responseEntity == null) { - classLogger.error("No response entity received"); - return null; - } - - try (BufferedReader reader = new BufferedReader( - new InputStreamReader(responseEntity.getContent(), StandardCharsets.UTF_8))) { - StringBuilder sb = new StringBuilder(); - String line; - while ((line = reader.readLine()) != null) { - sb.append(line); - } - - JSONObject kserveResponse = new JSONObject(sb.toString()); - - JSONObject modelResponse = KServeAdapter.formatKServeResponse(kserveResponse); - - return modelResponse; - } - - } - } catch (Exception e) { - classLogger.error("Error making HTTP request", e); - } - return null; - } - - @Override - public ModelTypeEnum getModelType() { - // TODO Auto-generated method stub - return null; - } - - @Override - public void close() throws IOException { - // TODO Auto-generated method stub - } - - /** - * - * @throws Exception - */ - private void checkModelUp() throws Exception { - RemoteModelStateEnum currentState = zkClient.getModelState(this.engineId); - // If not active after handling warming/cold states, return null - if (currentState != RemoteModelStateEnum.ACTIVE) { - initiateAndWaitForDeployment(300000); - } - } - - private String getModelUrl() throws Exception { - return this.createServiceUrl(Services.MODEL); - - } - - @Override - protected AskModelEngineResponse askCall(InputMessage inputMessage, Insight insight, String roomId, - Map hyperParameters) { - try { - checkModelUp(); - String modelUrl = getModelUrl(); - - if (modelUrl.endsWith("/infer")) { - String pattern = "/v2/models/[^/]+/infer$"; - modelUrl = modelUrl.replaceAll(pattern, "/v1"); - } - - classLogger.info("Adding cluster address to parameters: {}", modelUrl); - if (hyperParameters != null) { - hyperParameters.put("base_url", modelUrl); - } else { - hyperParameters = new HashMap<>(); - hyperParameters.put("base_url", modelUrl); - } - - hyperParameters.put("stream", false); - return implementingEngineClass.askCall(inputMessage, insight, roomId, hyperParameters); - } catch (Exception e) { - classLogger.error("Error getting model URL or deploying model", e); - return null; - } - } - - @Override - protected EmbeddingsModelEngineResponse embeddingsCall(List stringsToEmbed, Insight insight, - Map parameters) { - try { - checkModelUp(); - String modelUrl = getModelUrl(); - if (modelUrl.endsWith("/infer")) { - String pattern = "/v2/models/[^/]+/infer$"; - modelUrl = modelUrl.replaceAll(pattern, "/v1"); - } - classLogger.info("Adding cluster address to parameters: {}", modelUrl); - if (parameters != null) { - parameters.put("base_url", modelUrl); - } else { - parameters = new HashMap<>(); - parameters.put("base_url", modelUrl); - } - parameters.put("stream", false); - return implementingEngineClass.embeddingsCall(stringsToEmbed, insight, parameters); - } catch (Exception e) { - classLogger.error("Error getting model URL or deploying model", e); - return null; - } - } - -} diff --git a/src/prerna/engine/impl/model/KServeImageEmbedEngine.java b/src/prerna/engine/impl/model/KServeImageEmbedEngine.java deleted file mode 100644 index 6075fd42095..00000000000 --- a/src/prerna/engine/impl/model/KServeImageEmbedEngine.java +++ /dev/null @@ -1,81 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model; - -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import org.json.JSONObject; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.engine.api.ModelTypeEnum; -import prerna.engine.impl.model.responses.EmbeddingsModelEngineResponse; -import prerna.om.Insight; - -public class KServeImageEmbedEngine extends AbstractRemoteModelEngine { - - private static final Logger classLogger = LogManager.getLogger(KServeImageEmbedEngine.class); - - @Override - public EmbeddingsModelEngineResponse embeddingsCall(List imagesToEmbed, Insight insight, Map parameters) { - classLogger.debug("Handling KServeImageEmbed Request.."); - - JSONObject payload = new JSONObject(); - - payload.put("image", imagesToEmbed); - - if (parameters != null && parameters.containsKey("pooling_strategy")) { - String poolingStrategy = (String) parameters.get("pooling_strategy"); - payload.put("pooling_strategy", poolingStrategy); - } - - classLogger.debug("KServeVision embeddingsCall payload: {}", payload.toString(2)); - - try { - JSONObject modelResponse = makeModelRequest(payload); - if (modelResponse != null) { - return EmbeddingsModelEngineResponse.fromJson(modelResponse); - } else { - classLogger.error("Received null response from model"); - List> emptyEmbeddings = new ArrayList<>(); - return new EmbeddingsModelEngineResponse(emptyEmbeddings, 0, 0); - } - } catch (Exception e) { - classLogger.error("Error making model request", e); - List> emptyEmbeddings = new ArrayList<>(); - return new EmbeddingsModelEngineResponse(emptyEmbeddings, 0, 0); - } - } - - @Override - public ModelTypeEnum getModelType() { - return ModelTypeEnum.KSERVE_IMAGE_EMBED; - } - -} diff --git a/src/prerna/engine/impl/model/KServeImageEngine.java b/src/prerna/engine/impl/model/KServeImageEngine.java deleted file mode 100644 index 5b6b210faec..00000000000 --- a/src/prerna/engine/impl/model/KServeImageEngine.java +++ /dev/null @@ -1,110 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model; - -import java.util.HashMap; -import java.util.Map; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONObject; - -import prerna.engine.api.ModelTypeEnum; -import prerna.engine.impl.model.message.InputMessage; -import prerna.engine.impl.model.responses.AskImageModelEngineResponse; -import prerna.om.Insight; - -public class KServeImageEngine extends AbstractRemoteModelEngine { - - private static final Logger classLogger = LogManager.getLogger(KServeImageEngine.class); - - @Override - public AskImageModelEngineResponse askCall(InputMessage inputMessage, Insight insight, String roomId, - Map hyperParameters) { - classLogger.debug("Handling KServeImage Request.."); - - JSONObject payload = new JSONObject(); - - payload.put("prompt", inputMessage.getFullInputPrompt()); - - if (hyperParameters != null) { - if (hyperParameters.containsKey("negative_prompt")) { - String negativePrompt = (String) hyperParameters.get("negative_prompt"); - payload.put("negative_prompt", negativePrompt); - } - - if (hyperParameters.containsKey("height")) { - Object heightObj = hyperParameters.get("height"); - payload.put("height", heightObj.toString()); - } - - if (hyperParameters.containsKey("width")) { - Object widthObj = hyperParameters.get("width"); - payload.put("width", widthObj.toString()); - } - - if (hyperParameters.containsKey("num_inference_steps")) { - Object stepsObj = hyperParameters.get("num_inference_steps"); - payload.put("num_inference_steps", stepsObj.toString()); - } - - if (hyperParameters.containsKey("guidance_scale")) { - Object scaleObj = hyperParameters.get("guidance_scale"); - payload.put("guidance_scale", scaleObj.toString()); - } - - if (hyperParameters.containsKey("num_images")) { - Object numImagesObj = hyperParameters.get("num_images"); - payload.put("num_images", numImagesObj.toString()); - } - } - - try { - JSONObject modelResponse = makeModelRequest(payload); - if (modelResponse != null) { - return AskImageModelEngineResponse.getKServeImageResponse(modelResponse); - } else { - classLogger.error("Received null response from model"); - Map responseMap = new HashMap<>(); - responseMap.put("output", "Error creating image."); - return new AskImageModelEngineResponse(responseMap, 0, 0); - } - } catch (Exception e) { - classLogger.error("Error making model request", e); - Map responseMap = new HashMap<>(); - responseMap.put("output", "Error creating image."); - return new AskImageModelEngineResponse(responseMap, 0, 0); - } - } - - @Override - public ModelTypeEnum getModelType() { - return ModelTypeEnum.KSERVE_IMAGE; - } - -} diff --git a/src/prerna/engine/impl/model/KServeTTSEngine.java b/src/prerna/engine/impl/model/KServeTTSEngine.java deleted file mode 100644 index 6006b94150e..00000000000 --- a/src/prerna/engine/impl/model/KServeTTSEngine.java +++ /dev/null @@ -1,104 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model; - -import java.util.HashMap; -import java.util.Map; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONObject; - -import prerna.engine.api.ModelTypeEnum; -import prerna.engine.impl.model.message.InputMessage; -import prerna.engine.impl.model.responses.AskTTSModelEngineResponse; -import prerna.om.Insight; - -public class KServeTTSEngine extends AbstractRemoteModelEngine { - - private static final Logger classLogger = LogManager.getLogger(KServeTTSEngine.class); - - @Override - public AskTTSModelEngineResponse askCall(InputMessage inputMessage, Insight insight, String roomId, - Map parameters) { - String inputText = inputMessage.getFullInputPrompt(); - classLogger.debug("Handling KServeTTS Request for text: {}", inputText); - - JSONObject payload = new JSONObject(); - - payload.put("text", inputText); - - if (parameters != null) { - if (parameters.containsKey("voice")) { - String voice = (String) parameters.get("voice"); - payload.put("voice", voice); - classLogger.debug("Using voice: {}", voice); - } - - if (parameters.containsKey("speed")) { - String speed = (String) parameters.get("speed"); - payload.put("speed", speed); - classLogger.debug("Using speed: {}", speed); - } - - } - - try { - JSONObject modelResponse = makeModelRequest(payload); - if (modelResponse != null) { - AskTTSModelEngineResponse ttsResponse = AskTTSModelEngineResponse.getKServeTTSResponse(modelResponse); - - classLogger.debug("TTS response received - Format: {}, Duration: {}, Voice: {}", - ttsResponse.getAudioFormat(), ttsResponse.getDuration(), ttsResponse.getVoice()); - - if (ttsResponse.hasError()) { - classLogger.error("TTS model returned error: {}", ttsResponse.getErrorMessage()); - } - - return ttsResponse; - } else { - classLogger.error("Received null response from TTS model"); - Map errorResponse = new HashMap<>(); - errorResponse.put("status", "error"); - errorResponse.put("message", "Null response from TTS model"); - return new AskTTSModelEngineResponse(errorResponse, 0, 0); - } - } catch (Exception e) { - classLogger.error("Error processing TTS request: {}", e.getMessage(), e); - Map errorResponse = new HashMap<>(); - errorResponse.put("status", "error"); - errorResponse.put("message", "Error processing TTS request: " + e.getMessage()); - return new AskTTSModelEngineResponse(errorResponse, 0, 0); - } - } - - @Override - public ModelTypeEnum getModelType() { - return ModelTypeEnum.KSERVE_TTS; - } -} diff --git a/src/prerna/engine/impl/model/KServeVisionEngine.java b/src/prerna/engine/impl/model/KServeVisionEngine.java deleted file mode 100644 index d664aeb53ab..00000000000 --- a/src/prerna/engine/impl/model/KServeVisionEngine.java +++ /dev/null @@ -1,87 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model; - -import java.util.Map; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONObject; - -import prerna.engine.api.ModelTypeEnum; -import prerna.engine.impl.model.message.InputMessage; -import prerna.engine.impl.model.responses.AskStringModelEngineResponse; -import prerna.om.Insight; - -public class KServeVisionEngine extends AbstractRemoteModelEngine { - - private static final Logger classLogger = LogManager.getLogger(KServeVisionEngine.class); - - @Override - public AskStringModelEngineResponse askCall(InputMessage inputMessage, Insight insight, String roomId, - Map hyperParameters) { - classLogger.debug("Handling KServeVision Request.."); - - JSONObject payload = new JSONObject(); - - // Vision models require images in the payload.. - if (hyperParameters != null && hyperParameters.containsKey("image_url")) { - String imageUrl = (String) hyperParameters.get("image_url"); - payload.put("image", imageUrl); - } else { - classLogger.warn("No image_url found in hyperParameters"); - AskStringModelEngineResponse response = new AskStringModelEngineResponse( - "Please provide an image parameter.", 0, 0); - return response; - } - - payload.put("text", inputMessage.getFullInputPrompt()); - - classLogger.debug("KServeVision askCall payload: {}", payload.toString(2)); - - try { - JSONObject modelResponse = makeModelRequest(payload); - if (modelResponse != null) { - return AskStringModelEngineResponse.fromJson(modelResponse); - } else { - classLogger.error("Received null response from model"); - return new AskStringModelEngineResponse("Error processing image.", 0, 0); - } - } catch (Exception e) { - classLogger.error("Error making model request", e); - return new AskStringModelEngineResponse("An error occurred while processing the request: " + e.getMessage(), - 0, 10); - } - } - - @Override - public ModelTypeEnum getModelType() { - return ModelTypeEnum.KSERVE_VISION; - } - -} diff --git a/src/prerna/engine/impl/model/KubernetesModelScaler.java b/src/prerna/engine/impl/model/KubernetesModelScaler.java deleted file mode 100644 index a90b247d5ee..00000000000 --- a/src/prerna/engine/impl/model/KubernetesModelScaler.java +++ /dev/null @@ -1,305 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model; - -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -import org.apache.http.HttpEntity; -import org.apache.http.client.config.RequestConfig; -import org.apache.http.client.methods.CloseableHttpResponse; -import org.apache.http.client.methods.HttpGet; -import org.apache.http.client.methods.HttpPost; -import org.apache.http.entity.StringEntity; -import org.apache.http.impl.client.CloseableHttpClient; -import org.apache.http.impl.client.HttpClients; -import org.apache.http.util.EntityUtils; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONArray; -import org.json.JSONObject; - -import prerna.cluster.util.IRemoteClientServer; -import prerna.cluster.util.RemoteClientServerZKRESTProxy; -import prerna.cluster.util.ZKClientFactory; - -public final class KubernetesModelScaler { - - private static final Logger classLogger = LogManager.getLogger(KubernetesModelScaler.class); - - private static volatile KubernetesModelScaler instance; - - private KubernetesModelScaler() { - classLogger.info("KubernetesModelScaler being initialized..."); - } - - public static KubernetesModelScaler getInstance() { - if (instance != null) { - return instance; - } - - if (instance == null) { - synchronized (KubernetesModelScaler.class) { - if (instance == null) { - instance = new KubernetesModelScaler(); - instance.init(); - } - } - } - return instance; - } - - private IRemoteClientServer zkClient; - // Use this to simulate the cluster environment - private Boolean devPortForwarding = false; - // For normal development - private String kmsUrl = null; - - private void init() { - - // Get the appropriate ZK client implementation based on environment - this.zkClient = ZKClientFactory.getZKClient(this.devPortForwarding); - - // Check if we're using the REST proxy (for KMS_INGRESS validation) - boolean usingRestProxy = this.zkClient instanceof RemoteClientServerZKRESTProxy; - // Check if we have an ingress env var - String kmsIngressUrl = System.getenv("KMS_INGRESS"); - - if (kmsIngressUrl != null && !kmsIngressUrl.isEmpty()) { - classLogger.info("Using KMS_INGRESS from environment: {}", kmsIngressUrl); - if (kmsIngressUrl.endsWith("/")) { - kmsIngressUrl += "/"; - } - this.kmsUrl = kmsIngressUrl; - } else if (this.devPortForwarding) { - classLogger.info("Using devPortforwarding for KMS URL with localhost:8000/"); - this.kmsUrl = "http://localhost:8000/"; - } else { - classLogger.info( - "KMS_INGRESS environment variable not found and devPortforwarding not set, using ZooKeeper for KMS IP resolution. This is correct for production deployments."); - String zkModelScalerIp = zkClient.getModelScalerIp(); - - if (zkModelScalerIp == null || zkModelScalerIp.trim().isEmpty()) { - throw new RuntimeException( - "Unable to determine KMS URL: ZooKeeper returned null or empty model scaler IP"); - } - - if (!zkModelScalerIp.startsWith("http://") && !zkModelScalerIp.startsWith("https://")) { - zkModelScalerIp = "http://" + zkModelScalerIp; - } - - this.kmsUrl = zkModelScalerIp; - } - } - - public Map getNodePoolsInfo() throws Exception { - String serviceUrl = this.kmsUrl + "/api/resources/node-pools-info"; - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(5000).setSocketTimeout(5000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpGet httpGet = new HttpGet(serviceUrl); - httpGet.setHeader("Accept", "application/json"); - - try (CloseableHttpResponse response = httpClient.execute(httpGet)) { - int statusCode = response.getStatusLine().getStatusCode(); - HttpEntity responseEntity = response.getEntity(); - String responseString = EntityUtils.toString(responseEntity); - - if (statusCode == 200) { - JSONObject jsonResponse = new JSONObject(responseString); - - Map result = new HashMap<>(); - - if (jsonResponse.has("message")) { - result.put("message", jsonResponse.getString("message")); - } - - if (jsonResponse.has("zk_info")) { - JSONObject zkInfoObj = jsonResponse.getJSONObject("zk_info"); - Map zkInfoMap = jsonToMap(zkInfoObj); - result.put("zk_info", zkInfoMap); - } - - if (jsonResponse.has("pools")) { - JSONArray poolsArray = jsonResponse.getJSONArray("pools"); - List> poolsList = new ArrayList<>(); - - for (int i = 0; i < poolsArray.length(); i++) { - JSONObject poolObject = poolsArray.getJSONObject(i); - Map poolMap = jsonToMap(poolObject); - poolsList.add(poolMap); - } - - result.put("pools", poolsList); - } - - if (jsonResponse.has("active_models_actual")) { - JSONObject activeModelsObj = jsonResponse.getJSONObject("active_models_actual"); - Map activeModelsMap = jsonToMap(activeModelsObj); - result.put("active_models_actual", activeModelsMap); - } - - classLogger.info("Successfully retrieved node pool information"); - return result; - } else { - classLogger.error("Error retrieving node pool information: {} (Status: {})", responseString, - statusCode); - throw new RuntimeException("Failed to retrieve node pool information: " + responseString); - } - } - } catch (Exception e) { - classLogger.error("Error making request to model scaler for node pool info: {}", e.getMessage(), e); - throw new RuntimeException("Failed to retrieve node pool information", e); - } - } - - private Map jsonToMap(JSONObject json) { - Map map = new HashMap<>(); - - for (String key : json.keySet()) { - Object value = json.get(key); - - if (value instanceof JSONObject) { - map.put(key, jsonToMap((JSONObject) value)); - } else if (value instanceof JSONArray) { - JSONArray array = (JSONArray) value; - List list = new ArrayList<>(); - - for (int i = 0; i < array.length(); i++) { - Object item = array.get(i); - if (item instanceof JSONObject) { - list.add(jsonToMap((JSONObject) item)); - } else { - list.add(item); - } - } - - map.put(key, list); - } else { - map.put(key, value); - } - } - - return map; - } - - /** - * Retrieve the cached (or freshly re-loaded) KServe deployment configs from the - * model-scaler service. - * - * @param refresh if {@code true}, tells the micro-service to reload the configs - * from GCS before returning the response. - * @return a Map keyed by model-name, where each value is a nested Map - * representing the InferenceService spec for that model. - * @throws Exception on any HTTP or parse failure. - */ - public Map getModelDeploymentConfigs(boolean refresh) throws Exception { - // Build the URL – keep the same /api/resources prefix you use elsewhere - String endpoint = "/api/resources/model-deploy-configs"; - String serviceUrl = this.kmsUrl + endpoint + (refresh ? "?refresh=true" : ""); - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(5000).setSocketTimeout(5000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpGet httpGet = new HttpGet(serviceUrl); - httpGet.setHeader("Accept", "application/json"); - - try (CloseableHttpResponse response = httpClient.execute(httpGet)) { - int statusCode = response.getStatusLine().getStatusCode(); - String responseBody = EntityUtils.toString(response.getEntity()); - - if (statusCode == 200) { - JSONObject jsonResponse = new JSONObject(responseBody); - Map configs = jsonToMap(jsonResponse); - - classLogger.info("Retrieved {} model deployment configs (refresh={})", configs.size(), refresh); - return configs; - } else { - classLogger.error("Failed to fetch model deployment configs: {} (status {})", responseBody, - statusCode); - throw new RuntimeException("Failed to retrieve model deployment configs: " + responseBody); - } - } - } catch (Exception e) { - classLogger.error("Error contacting model-scaler for deployment configs: {}", e.getMessage(), e); - throw new RuntimeException("Failed to retrieve model deployment configs", e); - } - } - - public Map canItRun(String hfModelId) throws Exception { - - String serviceUrl = this.kmsUrl + "/api/can-it-run"; - - RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(5000).setSocketTimeout(5000).build(); - - try (CloseableHttpClient httpClient = HttpClients.custom().setDefaultRequestConfig(requestConfig).build()) { - - HttpPost httpPost = new HttpPost(serviceUrl); - httpPost.setHeader("Content-Type", "application/json"); - - JSONObject requestBody = new JSONObject(); - requestBody.put("model_id", hfModelId); - - StringEntity entity = new StringEntity(requestBody.toString()); - httpPost.setEntity(entity); - - try (CloseableHttpResponse response = httpClient.execute(httpPost)) { - int statusCode = response.getStatusLine().getStatusCode(); - HttpEntity responseEntity = response.getEntity(); - String responseString = EntityUtils.toString(responseEntity); - - if (statusCode == 200) { - JSONObject jsonResponse = new JSONObject(responseString); - - Map result = new HashMap<>(); - for (String key : jsonResponse.keySet()) { - result.put(key, jsonResponse.get(key)); - } - - classLogger.info("Successfully checked compatibility for model: {} - Can run: {}", hfModelId, - result.get("can_run")); - return result; - } else { - JSONObject errorResponse = new JSONObject(responseString); - String errorMessage = errorResponse.getJSONObject("detail").getString("message"); - classLogger.error("Error checking model compatibility: {} (Status: {})", errorMessage, statusCode); - throw new RuntimeException("Failed to check model compatibility: " + errorMessage); - } - } - } catch (Exception e) { - classLogger.error("Error making request to model scaler: {}", e.getMessage(), e); - throw new RuntimeException("Failed to check model compatibility", e); - } - } - -} diff --git a/src/prerna/engine/impl/model/NEREngine.java b/src/prerna/engine/impl/model/NEREngine.java deleted file mode 100644 index 4ea60d80f4e..00000000000 --- a/src/prerna/engine/impl/model/NEREngine.java +++ /dev/null @@ -1,95 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model; - -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import org.json.JSONObject; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.engine.impl.model.responses.NerModelEngineResponse; -import prerna.om.Insight; -import prerna.engine.api.ModelTypeEnum; - -/** - * Named Entity Recognition models - */ -public class NEREngine extends AbstractRemoteModelEngine { - - private static final Logger classLogger = LogManager.getLogger(NEREngine.class); - - public NerModelEngineResponse predict(String text, List entities, List maskEntities, Insight insight, Map parameters) { - JSONObject payload = new JSONObject(); - payload.put("text", text); - payload.put("labels", entities); - - if (maskEntities != null && !maskEntities.isEmpty()) { - payload.put("mask_entities", maskEntities); - } - - payload.put("model", this.model); - - classLogger.debug("NER predict payload: {}", payload.toString(2)); - - try { - JSONObject response = this.makeModelRequest(payload); - - if (response == null) { - Map errorMap = new HashMap<>(); - errorMap.put("status", "error"); - errorMap.put("message", "Null response from model"); - return new NerModelEngineResponse(errorMap, 0, 0); - } - - if (response.has("status") && "error".equals(response.getString("status"))) { - Map errorMap = new HashMap<>(); - errorMap.put("status", "error"); - errorMap.put("message", response.optString("message", "Unknown error")); - errorMap.put("code", response.optInt("code", 0)); - return new NerModelEngineResponse(errorMap, 0, 0); - } - - NerModelEngineResponse formattedResponse = NerModelEngineResponse.fromJson(response); - return formattedResponse; - } catch (Exception e) { - classLogger.error("Error making model request", e); - Map errorMap = new HashMap<>(); - errorMap.put("status", "error"); - errorMap.put("message", e.getMessage()); - - return new NerModelEngineResponse(errorMap, 0, 0); - } - } - - @Override - public ModelTypeEnum getModelType() { - return ModelTypeEnum.NER; - } -} \ No newline at end of file diff --git a/src/prerna/engine/impl/model/kserve/KServeAdapter.java b/src/prerna/engine/impl/model/kserve/KServeAdapter.java deleted file mode 100644 index 907c43b15d4..00000000000 --- a/src/prerna/engine/impl/model/kserve/KServeAdapter.java +++ /dev/null @@ -1,142 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model.kserve; - -import java.util.List; -import org.json.JSONObject; -import org.json.JSONArray; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -public class KServeAdapter { - /** - * Converts a model-specific payload to a KServe inference request - */ - public static JSONObject toKServeRequest(JSONObject modelPayload) { - Logger adapterLogger = LogManager.getLogger(KServeAdapter.class); - - JSONObject kserveRequest = new JSONObject(); - JSONArray inputs = new JSONArray(); - - adapterLogger.debug("Converting model payload to KServe format: {}", modelPayload.toString(2)); - - for (String key : modelPayload.keySet()) { - if (key.equals("model")) continue; - - JSONObject input = new JSONObject(); - input.put("name", key); - - Object value = modelPayload.get(key); - if (value instanceof String) { - input.put("shape", new int[]{1}); - input.put("datatype", "BYTES"); - JSONArray data = new JSONArray(); - data.put(value); - input.put("data", data); - } else if (value instanceof List) { - List list = (List) value; - input.put("shape", new int[]{list.size()}); - input.put("datatype", "BYTES"); - - JSONArray dataArray = new JSONArray(); - for (Object item : list) { - dataArray.put(item); - } - input.put("data", dataArray); - } else if (value instanceof JSONArray) { - JSONArray array = (JSONArray) value; - input.put("shape", new int[]{array.length()}); - input.put("datatype", "BYTES"); - input.put("data", array); - } - - inputs.put(input); - } - - kserveRequest.put("inputs", inputs); - - adapterLogger.debug("Final KServe request: {}", kserveRequest.toString(2)); - - return kserveRequest; - } - - /** - * Extracts data from KServe response - */ - public static JSONObject formatKServeResponse(JSONObject kserveResponse) { - Logger adapterLogger = LogManager.getLogger(KServeAdapter.class); - JSONObject result = new JSONObject(); - - adapterLogger.debug("Formatting KServe response: {}", kserveResponse.toString(2)); - - if (kserveResponse.has("outputs")) { - JSONArray outputs = kserveResponse.getJSONArray("outputs"); - - for (int i = 0; i < outputs.length(); i++) { - JSONObject output = outputs.getJSONObject(i); - String name = output.getString("name"); - JSONArray data = output.getJSONArray("data"); - - if (data.length() > 0) { - if (data.length() == 1) { - Object dataValue = data.get(0); - - if (dataValue instanceof String) { - String value = (String) dataValue; - - if ((value.startsWith("{") && value.endsWith("}")) || - (value.startsWith("[") && value.endsWith("]"))) { - try { - if (value.startsWith("{")) { - result.put(name, new JSONObject(value)); - } - else if (value.startsWith("[")) { - result.put(name, new JSONArray(value)); - } - } catch (Exception e) { - adapterLogger.debug("Could not parse JSON string for field {}: {}", name, e.getMessage()); - result.put(name, value); - } - } else { - result.put(name, value); - } - } else { - result.put(name, dataValue); - } - } else { - result.put(name, data); - } - } else { - result.put(name, new JSONArray()); - } - } - } - - return result; - } -} diff --git a/src/prerna/engine/impl/model/responses/AskImageModelEngineResponse.java b/src/prerna/engine/impl/model/responses/AskImageModelEngineResponse.java index b602444b6a2..fefb63f2439 100644 --- a/src/prerna/engine/impl/model/responses/AskImageModelEngineResponse.java +++ b/src/prerna/engine/impl/model/responses/AskImageModelEngineResponse.java @@ -31,19 +31,13 @@ import java.util.List; import java.util.Map; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; import org.json.JSONArray; -import org.json.JSONException; import org.json.JSONObject; public class AskImageModelEngineResponse extends AskModelEngineResponse> { - private static final Logger classLogger = LogManager.getLogger(AskImageModelEngineResponse.class); - private static final long serialVersionUID = 1L; - public static final String SOURCE_KSERVE = "kserve"; public static final String SOURCE_OPENAI = "openai"; /** @@ -71,7 +65,6 @@ public static AskImageModelEngineResponse getOpenAIImageResponse(List im Map responseMap = new HashMap<>(); if (imageList != null && !imageList.isEmpty()) { - // Convert List to String[] for consistency with KServe format String[] imageArray = imageList.toArray(new String[0]); responseMap.put("images", imageArray); responseMap.put("source", SOURCE_OPENAI); @@ -91,83 +84,6 @@ public static AskImageModelEngineResponse getOpenAIImageResponse(List im return new AskImageModelEngineResponse(responseMap, numberOfTokensInPrompt, numberOfTokensInResponse); } - /** - * Factory method for KServe image responses - * - * @param response - * @return - */ - public static AskImageModelEngineResponse getKServeImageResponse(JSONObject response) { - if (response != null) { - Map responseMap = new HashMap<>(); - responseMap.put("source", SOURCE_KSERVE); - - if (response.has("output")) { - Object outputObj = response.get("output"); - String outputStr = null; - - if (outputObj instanceof JSONArray) { - JSONArray images = (JSONArray) outputObj; - String[] imageList = new String[images.length()]; - for (int i = 0; i < images.length(); i++) { - imageList[i] = images.getString(i); - } - responseMap.put("images", imageList); - responseMap.put("count", imageList.length); - } else if (outputObj instanceof String) { - outputStr = (String) outputObj; - try { - JSONArray images = new JSONArray(outputStr); - String[] imageList = new String[images.length()]; - for (int i = 0; i < images.length(); i++) { - imageList[i] = images.getString(i); - } - responseMap.put("images", imageList); - responseMap.put("count", imageList.length); - } catch (JSONException e) { - responseMap.put("images", outputStr); - responseMap.put("count", 1); - } - } else { - responseMap.put("images", outputObj); - responseMap.put("count", 1); - } - } - - if (response.has("prompt")) { - responseMap.put("prompt", response.getString("prompt")); - } - if (response.has("negative_prompt")) { - responseMap.put("negative_prompt", response.getString("negative_prompt")); - } - if (response.has("height")) { - responseMap.put("height", response.getInt("height")); - } - if (response.has("width")) { - responseMap.put("width", response.getInt("width")); - } - if (response.has("num_inference_steps")) { - responseMap.put("num_inference_steps", response.getInt("num_inference_steps")); - } - if (response.has("guidance_scale")) { - responseMap.put("guidance_scale", response.getDouble("guidance_scale")); - } - if (response.has("seed")) { - responseMap.put("seed", response.getInt("seed")); - } - - return new AskImageModelEngineResponse(responseMap, 0, 0); - } else { - classLogger.error("Null response from model request"); - Map errorMap = new HashMap<>(); - errorMap.put("status", "error"); - errorMap.put("message", "Null response from model request"); - errorMap.put("source", SOURCE_KSERVE); - - return new AskImageModelEngineResponse(errorMap, 0, 0); - } - } - /** * Get the images as a String array, regardless of source * @@ -233,7 +149,7 @@ public Boolean isBase64Format() { } /** - * Get the source of the images (kserve, openai, etc.) + * Get the source of the images (openai, etc.) * * @return Source string */ diff --git a/src/prerna/engine/impl/model/responses/AskStringModelEngineResponse.java b/src/prerna/engine/impl/model/responses/AskStringModelEngineResponse.java index 5c53a5a467a..a21b6f0832f 100644 --- a/src/prerna/engine/impl/model/responses/AskStringModelEngineResponse.java +++ b/src/prerna/engine/impl/model/responses/AskStringModelEngineResponse.java @@ -27,8 +27,6 @@ *******************************************************************************/ package prerna.engine.impl.model.responses; -import org.json.JSONObject; - public class AskStringModelEngineResponse extends AskModelEngineResponse { private static final long serialVersionUID = 1L; @@ -42,44 +40,5 @@ public String getStringResponse() { return this.getResponse(); } - /** - * Creates an AskModelEngineResponse from a JSONObject returned by the KServe adapter - * - * @param jsonResponse The JSONObject from makeModelRequest - * @return AskModelEngineResponse constructed from the JSONObject - */ - public static AskStringModelEngineResponse fromJson(JSONObject jsonResponse) { - if (jsonResponse == null) { - return null; - } - - String responseText; - if (jsonResponse.has("output")) { - Object outputObj = jsonResponse.get("output"); - if (outputObj instanceof JSONObject || outputObj instanceof org.json.JSONArray) { - responseText = outputObj.toString(); - } else { - responseText = jsonResponse.getString("output"); - } - } else { - responseText = ""; - } - - Integer promptTokens = 0; - Integer responseTokens = 0; - - if (jsonResponse.has("input_tokens")) { - promptTokens = jsonResponse.getInt("input_tokens"); - } - - if (jsonResponse.has("output_tokens")) { - responseTokens = jsonResponse.getInt("output_tokens"); - } - - AskStringModelEngineResponse response = new AskStringModelEngineResponse(responseText, promptTokens, responseTokens); - - return response; - } - // Additional methods specific to string responses can be added here } \ No newline at end of file diff --git a/src/prerna/engine/impl/model/responses/AskTTSModelEngineResponse.java b/src/prerna/engine/impl/model/responses/AskTTSModelEngineResponse.java deleted file mode 100644 index eb516ba64a9..00000000000 --- a/src/prerna/engine/impl/model/responses/AskTTSModelEngineResponse.java +++ /dev/null @@ -1,278 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model.responses; - -import java.util.HashMap; -import java.util.Map; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONException; -import org.json.JSONObject; - -public class AskTTSModelEngineResponse extends AskModelEngineResponse> { - - private static final Logger classLogger = LogManager.getLogger(AskTTSModelEngineResponse.class); - private static final long serialVersionUID = 1L; - - public static final String SOURCE_KSERVE = "kserve"; - - /** - * Constructor for TTS model response - * @param response The response map containing TTS data - * @param numberOfTokensInPrompt Number of tokens in the input prompt - * @param numberOfTokensInResponse Number of tokens in the response - */ - public AskTTSModelEngineResponse(Map response, Integer numberOfTokensInPrompt, Integer numberOfTokensInResponse) { - super(response, numberOfTokensInPrompt, numberOfTokensInResponse); - this.messageType = TTS; - } - - /** - * Factory method for KServe TTS responses - * @param response JSONObject response from KServe TTS model - * @return AskTTSModelEngineResponse instance - */ - public static AskTTSModelEngineResponse getKServeTTSResponse(JSONObject response) { - if (response != null) { - Map responseMap = new HashMap<>(); - responseMap.put("source", SOURCE_KSERVE); - - if (response.has("output")) { - Object outputObj = response.get("output"); - - if (outputObj instanceof String) { - responseMap.put("audio_data", outputObj); - responseMap.put("format", "base64"); - } else if (outputObj instanceof JSONObject) { - JSONObject outputJson = (JSONObject) outputObj; - - if (outputJson.has("audio")) { - responseMap.put("audio_data", outputJson.get("audio")); - responseMap.put("format", "base64"); - } - - if (outputJson.has("duration")) { - responseMap.put("duration", outputJson.get("duration")); - } - if (outputJson.has("voice")) { - responseMap.put("voice", outputJson.get("voice")); - } - if (outputJson.has("sample_rate")) { - responseMap.put("sample_rate", outputJson.get("sample_rate")); - } - if (outputJson.has("speed")) { - responseMap.put("speed", outputJson.get("speed")); - } - } else { - responseMap.put("audio_data", outputObj); - responseMap.put("format", "unknown"); - } - } - - if (response.has("audio_format")) { - responseMap.put("audio_format", response.getString("audio_format")); - } else { - responseMap.put("audio_format", "wav"); - } - - if (response.has("sample_rate")) { - responseMap.put("sample_rate", response.getInt("sample_rate")); - } else { - responseMap.put("sample_rate", 22050); - } - - if (response.has("duration")) { - responseMap.put("duration", response.getDouble("duration")); - } - - if (response.has("voice")) { - responseMap.put("voice", response.getString("voice")); - } - - if (response.has("speed")) { - responseMap.put("speed", response.getString("speed")); - } - - if (response.has("text")) { - responseMap.put("text", response.getString("text")); - } - - if (response.has("error")) { - responseMap.put("error", response.getString("error")); - } - - return new AskTTSModelEngineResponse(responseMap, 0, 0); - } else { - classLogger.error("Null response from TTS model request"); - Map errorMap = new HashMap<>(); - errorMap.put("status", "error"); - errorMap.put("message", "Null response from TTS model request"); - errorMap.put("source", SOURCE_KSERVE); - - return new AskTTSModelEngineResponse(errorMap, 0, 0); - } - } - - /** - * Get the audio data (typically base64 encoded) - * @return Audio data as string, or null if not available - */ - public String getAudioData() { - Map response = this.getResponse(); - Object audioObj = response.get("audio_data"); - return (audioObj instanceof String) ? (String) audioObj : null; - } - - /** - * Get the audio format (wav, mp3, etc.) - * @return Audio format string - */ - public String getAudioFormat() { - Map response = this.getResponse(); - Object formatObj = response.get("audio_format"); - return (formatObj instanceof String) ? (String) formatObj : "wav"; - } - - /** - * Get the sample rate of the audio - * @return Sample rate as integer - */ - public int getSampleRate() { - Map response = this.getResponse(); - Object sampleRateObj = response.get("sample_rate"); - if (sampleRateObj instanceof Integer) { - return (Integer) sampleRateObj; - } - return 22050; - } - - /** - * Get the duration of the audio in seconds - * @return Duration as double, or -1 if not available - */ - public double getDuration() { - Map response = this.getResponse(); - Object durationObj = response.get("duration"); - if (durationObj instanceof Double) { - return (Double) durationObj; - } else if (durationObj instanceof Integer) { - return ((Integer) durationObj).doubleValue(); - } - return -1.0; - } - - /** - * Get the voice used for synthesis - * @return Voice name as string, or null if not specified - */ - public String getVoice() { - Map response = this.getResponse(); - Object voiceObj = response.get("voice"); - return (voiceObj instanceof String) ? (String) voiceObj : null; - } - - /** - * Get the speed used for synthesis - * @return Speed as string, or null if not specified - */ - public String getSpeed() { - Map response = this.getResponse(); - Object speedObj = response.get("speed"); - return (speedObj instanceof String) ? (String) speedObj : null; - } - - /** - * Get the original text that was synthesized - * @return Original text as string, or null if not available - */ - public String getOriginalText() { - Map response = this.getResponse(); - Object textObj = response.get("text"); - return (textObj instanceof String) ? (String) textObj : null; - } - - /** - * Check if the audio data is base64 encoded - * @return true if base64 format, false otherwise - */ - public boolean isBase64Format() { - Map response = this.getResponse(); - Object formatObj = response.get("format"); - return "base64".equals(formatObj); - } - - /** - * Get the source of the TTS response (kserve, etc.) - * @return Source string - */ - public String getSource() { - Map response = this.getResponse(); - Object sourceObj = response.get("source"); - return (sourceObj instanceof String) ? (String) sourceObj : "unknown"; - } - - /** - * Check if there was an error in the TTS generation - * @return true if there was an error, false otherwise - */ - public boolean hasError() { - Map response = this.getResponse(); - return response.containsKey("error") || "error".equals(response.get("status")); - } - - /** - * Get the error message if any - * @return Error message as string, or null if no error - */ - public String getErrorMessage() { - Map response = this.getResponse(); - Object errorObj = response.get("error"); - if (errorObj instanceof String) { - return (String) errorObj; - } - Object messageObj = response.get("message"); - return (messageObj instanceof String) ? (String) messageObj : null; - } - - @Override - public String getStringResponse() { - Map response = this.getResponse(); - JSONObject jsonObject = new JSONObject(); - - for (Map.Entry entry : response.entrySet()) { - try { - jsonObject.put(entry.getKey(), entry.getValue()); - } catch (JSONException e) { - classLogger.warn("Failed to add key '{}' to JSON response: {}", entry.getKey(), e.getMessage()); - } - } - - return jsonObject.toString(); - } -} \ No newline at end of file diff --git a/src/prerna/engine/impl/model/responses/NerModelEngineResponse.java b/src/prerna/engine/impl/model/responses/NerModelEngineResponse.java deleted file mode 100644 index b53cff1d988..00000000000 --- a/src/prerna/engine/impl/model/responses/NerModelEngineResponse.java +++ /dev/null @@ -1,152 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model.responses; - -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONArray; -import org.json.JSONObject; - -public class NerModelEngineResponse extends AbstractModelEngineResponse> { - - private static final Logger classLogger = LogManager.getLogger(NerModelEngineResponse.class); - private static final long serialVersionUID = 1L; - - public static final String MESSAGE_ID = "messageId"; - public static final String ROOM_ID = "roomId"; - - private String messageId; - private String roomId; - - /** - * - * @param response - * @param numberOfTokensInPrompt - * @param numberOfTokensInResponse - */ - public NerModelEngineResponse(Map response, Integer numberOfTokensInPrompt, Integer numberOfTokensInResponse) { - super(response, numberOfTokensInPrompt, numberOfTokensInResponse); - } - - public void setMessageId(String messageId) { - this.messageId = messageId; - } - - public String getMessageId() { - return this.messageId; - } - - public void setRoomId(String roomId) { - this.roomId = roomId; - } - - public String getRoomId() { - return this.roomId; - } - - /** - * - * @param response - * @return - */ - public static NerModelEngineResponse fromJson(JSONObject response) { - Logger classLogger = LogManager.getLogger(NerModelEngineResponse.class); - - if (response != null) { - Map responseMap = new HashMap<>(); - - if (response.has("mask_values")) { - JSONObject maskValues = response.getJSONObject("mask_values"); - Map maskValuesMap = new HashMap<>(); - for (String key : maskValues.keySet()) { - maskValuesMap.put(key, maskValues.getString(key)); - } - responseMap.put("mask_values", maskValuesMap); - } - - if (response.has("entities")) { - JSONArray entitiesArray = response.getJSONArray("entities"); - List entitiesList = new ArrayList<>(); - for (int i = 0; i < entitiesArray.length(); i++) { - entitiesList.add(entitiesArray.getString(i)); - } - responseMap.put("entities", entitiesList); - } - - if (response.has("raw_output")) { - JSONArray rawOutput = response.getJSONArray("raw_output"); - List> rawOutputList = new ArrayList<>(); - for (int i = 0; i < rawOutput.length(); i++) { - JSONObject entity = rawOutput.getJSONObject(i); - Map entityMap = new HashMap<>(); - entityMap.put("start", entity.getInt("start")); - entityMap.put("end", entity.getInt("end")); - entityMap.put("text", entity.getString("text")); - entityMap.put("label", entity.getString("label")); - entityMap.put("score", entity.getDouble("score")); - rawOutputList.add(entityMap); - } - responseMap.put("raw_output", rawOutputList); - } - - if (response.has("output")) { - responseMap.put("output", response.getString("output")); - } - - if (response.has("input")) { - responseMap.put("input", response.getString("input")); - } - - if (response.has("status")) { - responseMap.put("status", response.getString("status")); - } else { - responseMap.put("status", "success"); - } - - if (response.has("message")) { - responseMap.put("message", response.getString("message")); - } else { - responseMap.put("message", ""); - } - - return new NerModelEngineResponse(responseMap, 0, 0); - } else { - classLogger.error("Null response from model request"); - Map errorMap = new HashMap<>(); - errorMap.put("status", "error"); - errorMap.put("message", "Null response from model request"); - - return new NerModelEngineResponse(errorMap, 0, 0); - } - } -} diff --git a/src/prerna/reactor/automation/AutomationConstants.java b/src/prerna/reactor/automation/AutomationConstants.java index 52edced0928..c48511b3fff 100644 --- a/src/prerna/reactor/automation/AutomationConstants.java +++ b/src/prerna/reactor/automation/AutomationConstants.java @@ -200,7 +200,6 @@ private AutomationConstants() { public static final String NODE_MODEL_CHAT = "model.chat"; public static final String NODE_MODEL_EMBEDDINGS = "model.embeddings"; public static final String NODE_MODEL_VISION = "model.vision"; - public static final String NODE_MODEL_NER = "model.ner"; public static final String NODE_STORAGE_LIST = "storage.list"; public static final String NODE_STORAGE_READ = "storage.read"; public static final String NODE_STORAGE_UPLOAD = "storage.upload"; diff --git a/src/prerna/reactor/automation/definition/AutomationDefinitionValidator.java b/src/prerna/reactor/automation/definition/AutomationDefinitionValidator.java index cd7d5860d76..fa37cf552f5 100644 --- a/src/prerna/reactor/automation/definition/AutomationDefinitionValidator.java +++ b/src/prerna/reactor/automation/definition/AutomationDefinitionValidator.java @@ -255,11 +255,6 @@ private static void validateNodeConfig(String nodeId, AutomationNodeType nodeTyp requireConfigString(nodeId, config, "text"); validateOptionalConfigObject(nodeId, config, AutomationConstants.CONFIG_PARAM_VALUES); } - case MODEL_NER -> { - requireConfigString(nodeId, config, "text"); - requireConfigStringListOrPlaceholder(nodeId, config, "entities"); - validateOptionalStringListOrPlaceholder(nodeId, config, "maskEntities"); - } case MODEL_VISION -> { requireConfigString(nodeId, config, "prompt"); validateOptionalConfigString(nodeId, config, "image"); diff --git a/src/prerna/reactor/automation/definition/AutomationNodeCatalog.java b/src/prerna/reactor/automation/definition/AutomationNodeCatalog.java index 6bd1d462b3a..0b33c4e1850 100644 --- a/src/prerna/reactor/automation/definition/AutomationNodeCatalog.java +++ b/src/prerna/reactor/automation/definition/AutomationNodeCatalog.java @@ -140,13 +140,6 @@ private static List createDefinitions() { field("prompt", ConfigFieldType.TEXT, "Question about the media", true, ""), field("paramValues", ConfigFieldType.JSON, "Model parameters", false, Map.of())), controlInputs(), controlAndResultOutputs())); - definitions.add(definition(AutomationNodeType.MODEL_NER, "Extract entities", "Find named entities in text.", - orderedConfig("engineId", "", "text", "", "entities", List.of(), "maskEntities", List.of()), - List.of(engineField(IEngine.CATALOG_TYPE.MODEL), - field("text", ConfigFieldType.TEXT, "Text to analyze", true, ""), - field("entities", ConfigFieldType.STRING_LIST, "Entity types", true, List.of()), - field("maskEntities", ConfigFieldType.STRING_LIST, "Entity types to mask", false, List.of())), - controlInputs(), controlAndResultOutputs())); definitions.add(storageDefinition(AutomationNodeType.STORAGE_LIST, "List files", "List files in connected storage.", false, false)); diff --git a/src/prerna/reactor/automation/definition/AutomationNodeType.java b/src/prerna/reactor/automation/definition/AutomationNodeType.java index 28db9f2f3e4..95d82c907fa 100644 --- a/src/prerna/reactor/automation/definition/AutomationNodeType.java +++ b/src/prerna/reactor/automation/definition/AutomationNodeType.java @@ -59,8 +59,6 @@ public enum AutomationNodeType { Permission.VIEW, true, true), MODEL_VISION(AutomationConstants.NODE_MODEL_VISION, Category.MODEL, IEngine.CATALOG_TYPE.MODEL, Permission.VIEW, true, true), - MODEL_NER(AutomationConstants.NODE_MODEL_NER, Category.MODEL, IEngine.CATALOG_TYPE.MODEL, Permission.VIEW, true, - true), STORAGE_LIST(AutomationConstants.NODE_STORAGE_LIST, Category.STORAGE, IEngine.CATALOG_TYPE.STORAGE, Permission.VIEW, true, true), STORAGE_READ(AutomationConstants.NODE_STORAGE_READ, Category.STORAGE, IEngine.CATALOG_TYPE.STORAGE, Permission.EDIT, diff --git a/src/prerna/reactor/automation/run/AutomationPixelExpressionRenderer.java b/src/prerna/reactor/automation/run/AutomationPixelExpressionRenderer.java index 2a9f4ec1886..54452c0cd52 100644 --- a/src/prerna/reactor/automation/run/AutomationPixelExpressionRenderer.java +++ b/src/prerna/reactor/automation/run/AutomationPixelExpressionRenderer.java @@ -174,15 +174,6 @@ static String model(String nodeType, Map config, Map map - && "error".equalsIgnoreCase(String.valueOf(map.get("status")))) { - throw new IllegalStateException(String.valueOf(map.get("message"))); - } if (AutomationConstants.NODE_MODEL_CHAT.equals(nodeType) || AutomationConstants.NODE_MODEL_VISION.equals(nodeType)) { Map metadata = new LinkedHashMap<>(); diff --git a/src/prerna/reactor/automation/run/AutomationRunExecutionService.java b/src/prerna/reactor/automation/run/AutomationRunExecutionService.java index 8d040ca6a2d..5a82a29c79c 100644 --- a/src/prerna/reactor/automation/run/AutomationRunExecutionService.java +++ b/src/prerna/reactor/automation/run/AutomationRunExecutionService.java @@ -409,7 +409,7 @@ private static Handler handlerFor(Map node) { Handler.DATABASE_WRITE; case AutomationConstants.NODE_VECTOR_ADD, AutomationConstants.NODE_VECTOR_DELETE -> Handler.VECTOR_WRITE; case AutomationConstants.NODE_MODEL_CHAT, AutomationConstants.NODE_MODEL_EMBEDDINGS, - AutomationConstants.NODE_MODEL_VISION, AutomationConstants.NODE_MODEL_NER -> + AutomationConstants.NODE_MODEL_VISION -> Handler.MODEL; case AutomationConstants.NODE_STORAGE_READ, AutomationConstants.NODE_STORAGE_UPLOAD, AutomationConstants.NODE_STORAGE_DELETE -> diff --git a/src/prerna/reactor/model/CanItRunReactor.java b/src/prerna/reactor/model/CanItRunReactor.java deleted file mode 100644 index 4352c6e32bb..00000000000 --- a/src/prerna/reactor/model/CanItRunReactor.java +++ /dev/null @@ -1,69 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.reactor.model; - -import java.util.Map; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.reactor.AbstractReactor; -import prerna.sablecc2.om.ReactorKeysEnum; -import prerna.sablecc2.om.nounmeta.NounMetadata; -import prerna.sablecc2.om.PixelDataType; -import prerna.sablecc2.om.PixelOperationType; -import prerna.engine.impl.model.KubernetesModelScaler; - -public class CanItRunReactor extends AbstractReactor { - - private static final Logger classLogger = LogManager.getLogger(CanItRunReactor.class); - - public CanItRunReactor() { - this.keysToGet = new String[] { ReactorKeysEnum.HF_MODEL_ID.getKey()}; - this.keyRequired = new int[] { 1 }; - } - - @Override - public NounMetadata execute() { - organizeKeys(); - String hfModelId = this.keyValue.get(this.keysToGet[0]); - - final KubernetesModelScaler kmsServer; - kmsServer = KubernetesModelScaler.getInstance(); - - try { - Map canRun = kmsServer.canItRun(hfModelId); - return new NounMetadata(canRun, PixelDataType.MAP, PixelOperationType.OPERATION); - } catch (Exception e) { - classLogger.error("Error checking if model can run: " + hfModelId, e); - throw new RuntimeException("Failed to check if model can run: " + e.getMessage()); - } - - } - -} diff --git a/src/prerna/reactor/model/GetNodePoolsInfoReactor.java b/src/prerna/reactor/model/GetNodePoolsInfoReactor.java deleted file mode 100644 index 8023d6079c9..00000000000 --- a/src/prerna/reactor/model/GetNodePoolsInfoReactor.java +++ /dev/null @@ -1,63 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.reactor.model; - -import java.util.Map; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.reactor.AbstractReactor; -import prerna.sablecc2.om.nounmeta.NounMetadata; -import prerna.sablecc2.om.PixelDataType; -import prerna.sablecc2.om.PixelOperationType; -import prerna.engine.impl.model.KubernetesModelScaler; -import prerna.auth.utils.SecurityAdminUtils; - -public class GetNodePoolsInfoReactor extends AbstractReactor { - - private static final Logger classLogger = LogManager.getLogger(GetNodePoolsInfoReactor.class); - - @Override - public NounMetadata execute() { - if (!SecurityAdminUtils.userIsAdmin(this.insight.getUser())) { - throw new IllegalArgumentException("User does not have permission to query this endpoint."); - } - - final KubernetesModelScaler kmsServer; - kmsServer = KubernetesModelScaler.getInstance(); - - try { - Map nodePoolsInfo = kmsServer.getNodePoolsInfo(); - return new NounMetadata(nodePoolsInfo, PixelDataType.MAP, PixelOperationType.OPERATION); - } catch (Exception e) { - classLogger.error("Error connecting to the Kubernetes Model Scaler endpoint for Nodepool information.."); - throw new RuntimeException("Failed to connect to Kubernetes Model Scaler endpoint: " + e.getMessage()); - } - } -} diff --git a/src/prerna/reactor/model/GetRemoteModelDeployConfigsReactor.java b/src/prerna/reactor/model/GetRemoteModelDeployConfigsReactor.java deleted file mode 100644 index 8d96fc9839c..00000000000 --- a/src/prerna/reactor/model/GetRemoteModelDeployConfigsReactor.java +++ /dev/null @@ -1,88 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.reactor.model; - -import java.util.Map; -import java.util.List; - - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.reactor.AbstractReactor; -import prerna.sablecc2.om.nounmeta.NounMetadata; -import prerna.sablecc2.om.PixelDataType; -import prerna.sablecc2.om.PixelOperationType; -import prerna.sablecc2.om.GenRowStruct; -import prerna.engine.impl.model.KubernetesModelScaler; -import prerna.auth.utils.SecurityAdminUtils; - - - -public class GetRemoteModelDeployConfigsReactor extends AbstractReactor { - - private static final Logger classLogger = LogManager.getLogger(GetRemoteModelDeployConfigsReactor.class); - - public GetRemoteModelDeployConfigsReactor() { - this.keysToGet = new String[] {"refresh"}; - this.keyRequired = new int [] { 0 }; - } - - @Override - public NounMetadata execute() { - if (!SecurityAdminUtils.userIsAdmin(this.insight.getUser())) { - throw new IllegalArgumentException("User does not have permission to query this endpoint."); - } - - this.organizeKeys(); - Boolean refresh = this.getRefreshBool(); - - final KubernetesModelScaler kmsServer; - kmsServer = KubernetesModelScaler.getInstance(); - - try { - Map nodePoolsInfo = kmsServer.getModelDeploymentConfigs(refresh); - return new NounMetadata(nodePoolsInfo, PixelDataType.MAP, PixelOperationType.OPERATION); - } catch (Exception e) { - classLogger.error("Error connecting to the Kubernetes Model Scaler endpoint for model deployment configurations.."); - throw new RuntimeException("Failed to connect to Kubernetes Model Scaler endpoint: " + e.getMessage()); - } - } - - private boolean getRefreshBool() { - GenRowStruct boolGrs = this.store.getGenRowStruct("refresh"); - if (boolGrs != null) { - if (boolGrs.size() > 0) { - List val = boolGrs.getValuesOfType(PixelDataType.BOOLEAN); - return (boolean) val.get(0); - } - } - return false; - } - -} diff --git a/src/prerna/reactor/model/MyRemoteModelsStatus.java b/src/prerna/reactor/model/MyRemoteModelsStatus.java deleted file mode 100644 index edcf4685af8..00000000000 --- a/src/prerna/reactor/model/MyRemoteModelsStatus.java +++ /dev/null @@ -1,80 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.reactor.model; - -import java.util.ArrayList; -import java.util.List; -import java.util.HashMap; -import java.util.Map; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.auth.utils.SecurityEngineUtils; -import prerna.reactor.AbstractReactor; -import prerna.sablecc2.om.PixelDataType; -import prerna.sablecc2.om.PixelOperationType; -import prerna.sablecc2.om.nounmeta.NounMetadata; -import prerna.cluster.util.IRemoteClientServer; -import prerna.cluster.util.RemoteModelInfo; -import prerna.cluster.util.ZKClientFactory; - - -public class MyRemoteModelsStatus extends AbstractReactor { - - private static final Logger classLogger = LogManager.getLogger(MyRemoteModelsStatus.class); - - @Override - public NounMetadata execute() { - final IRemoteClientServer zkClient = ZKClientFactory.getZKClient(false); - - List activeModels = zkClient.getActiveModels(); - List warmingModels = zkClient.getWarmingModels(); - - List myActiveModels = new ArrayList<>(); - List myWarmingModels = new ArrayList<>(); - - for (RemoteModelInfo model : activeModels) { - if (SecurityEngineUtils.userCanViewEngine(this.insight.getUser(), model.getId())) { - myActiveModels.add(model); - } - } - - for (RemoteModelInfo model : warmingModels) { - if (SecurityEngineUtils.userCanViewEngine(this.insight.getUser(), model.getId())) { - myWarmingModels.add(model); - } - } - - Map> modelsMap = new HashMap<>(); - modelsMap.put("activeModels", myActiveModels); - modelsMap.put("warmingModels", myWarmingModels); - - return new NounMetadata(modelsMap, PixelDataType.MAP, PixelOperationType.OPERATION); - } -} \ No newline at end of file diff --git a/src/prerna/reactor/model/NERReactor.java b/src/prerna/reactor/model/NERReactor.java deleted file mode 100644 index 1c42e2fa162..00000000000 --- a/src/prerna/reactor/model/NERReactor.java +++ /dev/null @@ -1,117 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.reactor.model; - -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Vector; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.auth.utils.SecurityEngineUtils; -import prerna.engine.api.IModelEngine; -import prerna.engine.impl.model.NEREngine; -import prerna.engine.impl.model.responses.NerModelEngineResponse; -import prerna.reactor.AbstractReactor; -import prerna.sablecc2.om.GenRowStruct; -import prerna.sablecc2.om.PixelDataType; -import prerna.sablecc2.om.PixelOperationType; -import prerna.sablecc2.om.ReactorKeysEnum; -import prerna.sablecc2.om.nounmeta.NounMetadata; -import prerna.util.Utility; - -public class NERReactor extends AbstractReactor { - - private static final Logger classLogger = LogManager.getLogger(NERReactor.class); - - public NERReactor() { - this.keysToGet = new String[] { ReactorKeysEnum.ENGINE.getKey(), ReactorKeysEnum.PROMPT.getKey(), - ReactorKeysEnum.ENTITIES.getKey(), ReactorKeysEnum.MASK_ENTITIES.getKey(), - ReactorKeysEnum.PARAM_VALUES_MAP.getKey() }; - this.keyRequired = new int[] { 1, 1, 1, 0, 0 }; - } - - @Override - public NounMetadata execute() { - organizeKeys(); - String engineId = this.keyValue.get(this.keysToGet[0]); - - if (!SecurityEngineUtils.userCanViewEngine(this.insight.getUser(), engineId)) { - throw new IllegalArgumentException( - "Model " + engineId + " does not exist or user does not have access to this model"); - } - - String prompt = this.keyValue.get(this.keysToGet[1]); - List entities = this.getListInput("entities"); - List maskEntities = this.getListInput("maskEntities"); - - Map paramMap = getMap(); - if (paramMap == null) { - paramMap = new HashMap(); - } - - // CASTING TO CORRECT ENGINE.. NER is not abstracted - IModelEngine targetModel = Utility.getModel(engineId); - NEREngine targetEngine = (NEREngine) targetModel; - - NerModelEngineResponse output = targetEngine.predict(prompt, entities, maskEntities, this.insight, paramMap); - - return new NounMetadata(output, PixelDataType.MAP, PixelOperationType.OPERATION); - - } - - private Map getMap() { - GenRowStruct mapGrs = this.store.getGenRowStruct(keysToGet[4]); - if (mapGrs != null && !mapGrs.isEmpty()) { - List mapInputs = mapGrs.getNounsOfType(PixelDataType.MAP); - if (mapInputs != null && !mapInputs.isEmpty()) { - return (Map) mapInputs.get(0).getValue(); - } - } - List mapInputs = this.curRow.getNounsOfType(PixelDataType.MAP); - if (mapInputs != null && !mapInputs.isEmpty()) { - return (Map) mapInputs.get(0).getValue(); - } - return null; - } - - private List getListInput(String noun) { - List colInputs = new Vector(); - GenRowStruct colGRS = this.store.getGenRowStruct(noun); - if (colGRS != null) { - for (int i = 0; i < colGRS.size(); i++) { - String stringValue = colGRS.get(i).toString(); - colInputs.add(stringValue); - } - } - return colInputs; - } - -} diff --git a/src/prerna/reactor/model/RemoteModelShutdownReactor.java b/src/prerna/reactor/model/RemoteModelShutdownReactor.java deleted file mode 100644 index 90a1b295ecd..00000000000 --- a/src/prerna/reactor/model/RemoteModelShutdownReactor.java +++ /dev/null @@ -1,73 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.reactor.model; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.auth.utils.SecurityEngineUtils; -import prerna.engine.api.IModelEngine; -import prerna.reactor.AbstractReactor; -import prerna.sablecc2.om.PixelDataType; -import prerna.sablecc2.om.PixelOperationType; -import prerna.sablecc2.om.ReactorKeysEnum; -import prerna.sablecc2.om.nounmeta.NounMetadata; -import prerna.util.Utility; -import prerna.engine.impl.model.AbstractRemoteModelEngine; - -public class RemoteModelShutdownReactor extends AbstractReactor { - - private static final Logger classLogger = LogManager.getLogger(RemoteModelShutdownReactor.class); - - public RemoteModelShutdownReactor() { - this.keysToGet = new String[] { ReactorKeysEnum.ENGINE.getKey()}; - this.keyRequired = new int[] { 1 }; - } - - @Override - public NounMetadata execute() { - organizeKeys(); - String engineId = this.keyValue.get(this.keysToGet[0]); - - if(!SecurityEngineUtils.userIsOwner(this.insight.getUser(), engineId)) { - throw new IllegalArgumentException("Only the owner can shutdown the model"); - } - - try { - IModelEngine targetModel = Utility.getModel(engineId); - AbstractRemoteModelEngine targetEngine = (AbstractRemoteModelEngine) targetModel; - - String shutdownResult = targetEngine.shutdownModelRequest(); - - return new NounMetadata(shutdownResult, PixelDataType.CONST_STRING, PixelOperationType.OPERATION); - } catch (Exception e) { - classLogger.error("Error shutting down model: " + engineId, e); - throw new RuntimeException("Failed to shutdown model: " + e.getMessage()); - } - } -} diff --git a/src/prerna/reactor/model/RemoteModelStartReactor.java b/src/prerna/reactor/model/RemoteModelStartReactor.java deleted file mode 100644 index 91a5a5b17ff..00000000000 --- a/src/prerna/reactor/model/RemoteModelStartReactor.java +++ /dev/null @@ -1,83 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.reactor.model; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import prerna.auth.utils.SecurityEngineUtils; -import prerna.engine.api.IModelEngine; -import prerna.reactor.AbstractReactor; -import prerna.sablecc2.om.PixelDataType; -import prerna.sablecc2.om.PixelOperationType; -import prerna.sablecc2.om.ReactorKeysEnum; -import prerna.sablecc2.om.nounmeta.NounMetadata; -import prerna.util.Utility; -import prerna.engine.impl.model.AbstractRemoteModelEngine; - -public class RemoteModelStartReactor extends AbstractReactor { - - private static final Logger classLogger = LogManager.getLogger(RemoteModelStartReactor.class); - - public RemoteModelStartReactor() { - this.keysToGet = new String[] { ReactorKeysEnum.ENGINE.getKey()}; - this.keyRequired = new int[] { 1 }; - } - - @Override - public NounMetadata execute() { - organizeKeys(); - String engineId = this.keyValue.get(this.keysToGet[0]); - - if(!SecurityEngineUtils.userCanViewEngine(this.insight.getUser(), engineId)) { - throw new IllegalArgumentException("Model " + engineId + " does not exist or user does not have access to this model"); - } - - try { - IModelEngine targetModel = Utility.getModel(engineId); - AbstractRemoteModelEngine targetEngine = (AbstractRemoteModelEngine) targetModel; - - Boolean startUpResult = targetEngine.initiateAndWaitForDeployment(200000); - - String result; - - if (startUpResult) { - result = "Model started successfully"; - } else { - result = "Model failed to start"; - } - - - return new NounMetadata(result, PixelDataType.CONST_STRING, PixelOperationType.OPERATION); - } catch (Exception e) { - classLogger.error("Error starting model: " + engineId, e); - throw new RuntimeException("Failed to start model: " + e.getMessage()); - } - } - -} diff --git a/src/prerna/sablecc2/om/ReactorKeysEnum.java b/src/prerna/sablecc2/om/ReactorKeysEnum.java index 5dde68b44ad..1fae4222bbb 100644 --- a/src/prerna/sablecc2/om/ReactorKeysEnum.java +++ b/src/prerna/sablecc2/om/ReactorKeysEnum.java @@ -97,7 +97,6 @@ public enum ReactorKeysEnum { END_ROW("endRow", "The ending row number to import from a file."), ENGINE("engine", "Id/Name of the engine"), ENGINE_TYPE("engineTypes", "The type of engine to filter into (DATABASE, STORAGE, MODEL, etc.)"), - ENTITIES("entities", "The entities to locate within a given text"), EXPIRES_ON("expires", "The date when the api expires"), EXPORT_AUDIT("exportAudit", "Boolean to include an audit sheet for parameters to excel exports"), EXPORT_TEMPLATE("export_template", "Template file to use for the export of this data"), @@ -134,7 +133,6 @@ public enum ReactorKeysEnum { HEADERS("headers", "All the headers strings we want to push as part of the header on the excel / powerpoint export"), HEADERS_MAP("headersMap", "Map containing key-value pairs to send in the headers of a request"), HEIGHT("height", "The height to use for screenshot capture."), - HF_MODEL_ID("hfModelId", "The model id for the HF model"), HOST("host", "The host URL"), HTML("html", "HTML text"), MARKDOWN("markdown", "Markdown text"), @@ -164,7 +162,6 @@ public enum ReactorKeysEnum { LOGICAL_NAME("logicalNames", "Column alias to be added to the master database"), MAINTAIN_COLUMNS("maintainCols", "Any additional columns to keep in the frame"), MAP("map", "Map that is the equivalent of a JSON for key-value properties"), - MASK_ENTITIES("maskEntities", "The entities to mask when returning results from an NER model"), MAX("max", "Maximum value of something. Typically a threshold"), MCP("mcp", "List of MCPs for the reactor to use"), MCP_METADATA("mcpMetadata", """ diff --git a/src/prerna/util/Settings.java b/src/prerna/util/Settings.java index 6687377cd24..f9fca93ac33 100644 --- a/src/prerna/util/Settings.java +++ b/src/prerna/util/Settings.java @@ -86,11 +86,8 @@ public class Settings { // FOR REMOTE CLIENTS - // The short name for a model.. Can't use the HuggingFace model repo id - // here due to naming restrictions on kubernetes resources + // The short name for a model public static final String MODEL = "MODEL"; - // The HuggingFace model repo id - public static final String MODEL_REPO_ID = "MODEL_REPO_ID"; // The type of model public static final String MODEL_TYPE = "MODEL_TYPE"; diff --git a/test/prerna/cluster/util/ModelAvailableListenerUnitTests.java b/test/prerna/cluster/util/ModelAvailableListenerUnitTests.java deleted file mode 100644 index 1362c85b8f7..00000000000 --- a/test/prerna/cluster/util/ModelAvailableListenerUnitTests.java +++ /dev/null @@ -1,173 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -import java.util.List; - -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.ZooKeeper; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -class ModelAvailableListenerUnitTests { - - private static final String TEST_PATH = "/models/testModel/status"; - private ModelZKServer mockServer; - private ZooKeeper mockZk; - private ModelAvailableListener listener; - - @BeforeEach - void setUp() { - mockServer = mock(ModelZKServer.class); - mockZk = mock(ZooKeeper.class); - listener = new ModelAvailableListener(TEST_PATH, mockServer); - } - - @Test - void testConstructorSetsPathAndServer() { - assertEquals(TEST_PATH, listener.getPath()); - assertNotNull(listener); - } - - @Test - void testGetPathReturnsPath() { - assertEquals(TEST_PATH, listener.getPath()); - } - - @Test - void testGetPathWithDifferentPath() { - ModelAvailableListener other = new ModelAvailableListener("/other/model/status", mockServer); - assertEquals("/other/model/status", other.getPath()); - } - - @Test - void testGetEventsReturnsNodeDataChanged() { - List events = listener.getEvents(); - assertNotNull(events); - assertEquals(1, events.size()); - assertEquals(EventType.NodeDataChanged, events.get(0)); - } - - @Test - void testGetPredicatesReturnsEquals() { - List predicates = listener.getPredicates(); - assertNotNull(predicates); - assertEquals(1, predicates.size()); - assertEquals("equals", predicates.get(0)); - } - - @Test - void testSetModelZKUpdatesServer() { - ModelZKServer newServer = mock(ModelZKServer.class); - assertDoesNotThrow(() -> listener.setModelZK(newServer)); - } - - @Test - void testAvailableConstant() { - assertEquals("Available", ModelAvailableListener.AVAILABLE); - } - - @Test - void testProcessWithAvailableStatusCallsGetNodeDataOnEndpoint() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("Available"); - - String expectedEndpointNode = TEST_PATH.replace("status", "endpoint"); - when(mockServer.getNodeData(expectedEndpointNode)).thenReturn("http://localhost:8080"); - - listener.process(TEST_PATH, mockZk); - - verify(mockServer).getNodeData(TEST_PATH); - verify(mockServer).getNodeData(expectedEndpointNode); - } - - @Test - void testProcessWithAvailableStatusCaseInsensitive() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("available"); - - String expectedEndpointNode = TEST_PATH.replace("status", "endpoint"); - when(mockServer.getNodeData(expectedEndpointNode)).thenReturn("http://localhost:9090"); - - listener.process(TEST_PATH, mockZk); - - verify(mockServer).getNodeData(TEST_PATH); - verify(mockServer).getNodeData(expectedEndpointNode); - } - - @Test - void testProcessWithFailStatusDoesNotCrash() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("Fail"); - - assertDoesNotThrow(() -> listener.process(TEST_PATH, mockZk)); - verify(mockServer).getNodeData(TEST_PATH); - } - - @Test - void testProcessWithFailStatusCaseInsensitive() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("fail"); - - assertDoesNotThrow(() -> listener.process(TEST_PATH, mockZk)); - verify(mockServer).getNodeData(TEST_PATH); - } - - @Test - void testProcessWithOtherStatusDoesNothingExtra() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("SomeOtherStatus"); - - listener.process(TEST_PATH, mockZk); - - // Only the initial getNodeData for the status should be called - verify(mockServer, times(1)).getNodeData(TEST_PATH); - verifyNoMoreInteractions(mockServer); - } - - @Test - void testProcessWithInitStatusDoesNothingExtra() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("INIT"); - - listener.process(TEST_PATH, mockZk); - - verify(mockServer, times(1)).getNodeData(TEST_PATH); - verifyNoMoreInteractions(mockServer); - } - - @Test - void testProcessEndpointNodeReplacesStatusWithEndpoint() { - String statusPath = "/models/myModel/status"; - ModelAvailableListener avListener = new ModelAvailableListener(statusPath, mockServer); - - when(mockServer.getNodeData(statusPath)).thenReturn("Available"); - when(mockServer.getNodeData("/models/myModel/endpoint")).thenReturn("http://host:1234"); - - avListener.process(statusPath, mockZk); - - verify(mockServer).getNodeData("/models/myModel/endpoint"); - } -} diff --git a/test/prerna/cluster/util/ModelFailListenerUnitTests.java b/test/prerna/cluster/util/ModelFailListenerUnitTests.java deleted file mode 100644 index 3929c18c0d2..00000000000 --- a/test/prerna/cluster/util/ModelFailListenerUnitTests.java +++ /dev/null @@ -1,159 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -import java.util.List; - -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.ZooKeeper; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -class ModelFailListenerUnitTests { - - private static final String TEST_PATH = "/models/testModel/fail"; - private ModelZKServer mockServer; - private ZooKeeper mockZk; - private ModelFailListener listener; - - @BeforeEach - void setUp() { - mockServer = mock(ModelZKServer.class); - mockZk = mock(ZooKeeper.class); - listener = new ModelFailListener(TEST_PATH, mockServer); - } - - @Test - void testConstructorSetsPathAndServer() { - assertEquals(TEST_PATH, listener.getPath()); - assertNotNull(listener); - } - - @Test - void testGetPathReturnsPath() { - assertEquals(TEST_PATH, listener.getPath()); - } - - @Test - void testGetPathWithDifferentPath() { - ModelFailListener other = new ModelFailListener("/other/path/fail", mockServer); - assertEquals("/other/path/fail", other.getPath()); - } - - @Test - void testGetEventsReturnsNodeChildrenChanged() { - List events = listener.getEvents(); - assertNotNull(events); - assertEquals(1, events.size()); - assertEquals(EventType.NodeChildrenChanged, events.get(0)); - } - - @Test - void testGetPredicatesReturnsEquals() { - List predicates = listener.getPredicates(); - assertNotNull(predicates); - assertEquals(1, predicates.size()); - assertEquals("equals", predicates.get(0)); - } - - @Test - void testSetModelZKUpdatesServer() { - ModelZKServer newServer = mock(ModelZKServer.class); - assertDoesNotThrow(() -> listener.setModelZK(newServer)); - } - - @Test - void testProcessWhenFailCountEqualsServerCountCallsUpdateNodeData() throws Exception { - List failedChildren = List.of("server1", "server2", "server3"); - List serverChildren = List.of("serverA", "serverB", "serverC"); - - when(mockZk.getChildren(TEST_PATH, true)).thenReturn(failedChildren); - when(mockZk.getChildren("/server", true)).thenReturn(serverChildren); - - listener.process(TEST_PATH, mockZk); - - String expectedStatusNode = TEST_PATH.replace("fail", "status"); - verify(mockServer).updateNodeData(expectedStatusNode, "FAIL", true); - } - - @Test - void testProcessWhenFailCountNotEqualServerCountDoesNotCallUpdateNodeData() throws Exception { - List failedChildren = List.of("server1"); - List serverChildren = List.of("serverA", "serverB", "serverC"); - - when(mockZk.getChildren(TEST_PATH, true)).thenReturn(failedChildren); - when(mockZk.getChildren("/server", true)).thenReturn(serverChildren); - - listener.process(TEST_PATH, mockZk); - - verify(mockServer, never()).updateNodeData(anyString(), anyString(), anyBoolean()); - } - - @Test - void testProcessStatusNodeReplacesFailWithStatus() throws Exception { - String pathWithFail = "/models/myModel/fail"; - ModelFailListener failListener = new ModelFailListener(pathWithFail, mockServer); - - List children = List.of("s1", "s2"); - when(mockZk.getChildren(pathWithFail, true)).thenReturn(children); - when(mockZk.getChildren("/server", true)).thenReturn(children); - - failListener.process(pathWithFail, mockZk); - - verify(mockServer).updateNodeData("/models/myModel/status", "FAIL", true); - } - - @Test - void testProcessWithEmptyChildrenLists() throws Exception { - List emptyList = List.of(); - - when(mockZk.getChildren(TEST_PATH, true)).thenReturn(emptyList); - when(mockZk.getChildren("/server", true)).thenReturn(emptyList); - - listener.process(TEST_PATH, mockZk); - - String expectedStatusNode = TEST_PATH.replace("fail", "status"); - verify(mockServer).updateNodeData(expectedStatusNode, "FAIL", true); - } - - @Test - void testProcessWithMoreFailedThanServers() throws Exception { - List failedChildren = List.of("s1", "s2", "s3"); - List serverChildren = List.of("s1", "s2"); - - when(mockZk.getChildren(TEST_PATH, true)).thenReturn(failedChildren); - when(mockZk.getChildren("/server", true)).thenReturn(serverChildren); - - listener.process(TEST_PATH, mockZk); - - verify(mockServer, never()).updateNodeData(anyString(), anyString(), anyBoolean()); - } -} diff --git a/test/prerna/cluster/util/ModelInitListenerUnitTests.java b/test/prerna/cluster/util/ModelInitListenerUnitTests.java deleted file mode 100644 index bc4a2c97ab0..00000000000 --- a/test/prerna/cluster/util/ModelInitListenerUnitTests.java +++ /dev/null @@ -1,181 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -import java.util.List; - -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.ZooKeeper; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -class ModelInitListenerUnitTests { - - private static final String TEST_PATH = "/models/testModel/status"; - private ModelZKServer mockServer; - private ZooKeeper mockZk; - private ModelInitListener listener; - - @BeforeEach - void setUp() { - mockServer = mock(ModelZKServer.class); - mockZk = mock(ZooKeeper.class); - listener = new ModelInitListener(TEST_PATH, mockServer); - } - - @Test - void testConstructorSetsPathAndServer() { - assertEquals(TEST_PATH, listener.getPath()); - assertNotNull(listener); - } - - @Test - void testGetPathReturnsPath() { - assertEquals(TEST_PATH, listener.getPath()); - } - - @Test - void testGetPathWithDifferentPath() { - ModelInitListener other = new ModelInitListener("/other/model/status", mockServer); - assertEquals("/other/model/status", other.getPath()); - } - - @Test - void testGetEventsReturnsNodeDataChanged() { - List events = listener.getEvents(); - assertNotNull(events); - assertEquals(1, events.size()); - assertEquals(EventType.NodeDataChanged, events.get(0)); - } - - @Test - void testGetPredicatesReturnsEquals() { - List predicates = listener.getPredicates(); - assertNotNull(predicates); - assertEquals(1, predicates.size()); - assertEquals("equals", predicates.get(0)); - } - - @Test - void testSetModelZKUpdatesServer() { - ModelZKServer newServer = mock(ModelZKServer.class); - assertDoesNotThrow(() -> listener.setModelZK(newServer)); - } - - @Test - void testProcessWithInitStatusAndCatchupFalseCallsSpinModel() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("INIT"); - - String modelNode = TEST_PATH.replace("/status", ""); - String smssValue = "model-smss-content"; - when(mockServer.getNodeData(modelNode)).thenReturn(smssValue); - - mockServer.catchup = false; - - listener.process(TEST_PATH, mockZk); - - verify(mockServer).spinModel(smssValue); - } - - @Test - void testProcessWithInitStatusCaseInsensitiveAndCatchupFalse() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("init"); - - String modelNode = TEST_PATH.replace("/status", ""); - String smssValue = "model-config"; - when(mockServer.getNodeData(modelNode)).thenReturn(smssValue); - - mockServer.catchup = false; - - listener.process(TEST_PATH, mockZk); - - verify(mockServer).spinModel(smssValue); - } - - @Test - void testProcessWithInitStatusAndCatchupTrueDoesNotCallSpinModel() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("INIT"); - - String modelNode = TEST_PATH.replace("/status", ""); - when(mockServer.getNodeData(modelNode)).thenReturn("model-smss"); - - mockServer.catchup = true; - - listener.process(TEST_PATH, mockZk); - - verify(mockServer, never()).spinModel(anyString()); - } - - @Test - void testProcessWithNonInitStatusDoesNotCallSpinModel() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("Available"); - - listener.process(TEST_PATH, mockZk); - - verify(mockServer, never()).spinModel(anyString()); - } - - @Test - void testProcessWithFailStatusDoesNotCallSpinModel() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("FAIL"); - - listener.process(TEST_PATH, mockZk); - - verify(mockServer, never()).spinModel(anyString()); - } - - @Test - void testProcessModelNodeIsPathWithoutStatus() { - String pathWithStatus = "/cluster/models/myEngine/status"; - ModelInitListener initListener = new ModelInitListener(pathWithStatus, mockServer); - - when(mockServer.getNodeData(pathWithStatus)).thenReturn("INIT"); - when(mockServer.getNodeData("/cluster/models/myEngine")).thenReturn("my-smss"); - - mockServer.catchup = false; - - initListener.process(pathWithStatus, mockZk); - - verify(mockServer).getNodeData("/cluster/models/myEngine"); - verify(mockServer).spinModel("my-smss"); - } - - @Test - void testProcessWithRandomStatusAndCatchupFalseDoesNotCallSpinModel() { - when(mockServer.getNodeData(TEST_PATH)).thenReturn("SomethingElse"); - - mockServer.catchup = false; - - listener.process(TEST_PATH, mockZk); - - verify(mockServer, never()).spinModel(anyString()); - } -} diff --git a/test/prerna/cluster/util/ModelUnavailableListenerUnitTests.java b/test/prerna/cluster/util/ModelUnavailableListenerUnitTests.java deleted file mode 100644 index 4efe80d07e5..00000000000 --- a/test/prerna/cluster/util/ModelUnavailableListenerUnitTests.java +++ /dev/null @@ -1,141 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -import java.util.List; - -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.ZooKeeper; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -class ModelUnavailableListenerUnitTests { - - private static final String TEST_PATH = "/models/testModel/endpoint"; - private ModelZKServer mockServer; - private ZooKeeper mockZk; - private ModelUnavailableListener listener; - - @BeforeEach - void setUp() { - mockServer = mock(ModelZKServer.class); - mockZk = mock(ZooKeeper.class); - listener = new ModelUnavailableListener(TEST_PATH, mockServer); - } - - @Test - void testConstructorSetsPathAndServer() { - assertEquals(TEST_PATH, listener.getPath()); - assertNotNull(listener); - } - - @Test - void testGetPathReturnsPath() { - assertEquals(TEST_PATH, listener.getPath()); - } - - @Test - void testGetPathWithDifferentPath() { - ModelUnavailableListener other = new ModelUnavailableListener("/other/endpoint", mockServer); - assertEquals("/other/endpoint", other.getPath()); - } - - @Test - void testGetEventsReturnsNodeDeleted() { - List events = listener.getEvents(); - assertNotNull(events); - assertEquals(1, events.size()); - assertEquals(EventType.NodeDeleted, events.get(0)); - } - - @Test - void testGetPredicatesReturnsEquals() { - List predicates = listener.getPredicates(); - assertNotNull(predicates); - assertEquals(1, predicates.size()); - assertEquals("equals", predicates.get(0)); - } - - @Test - void testSetModelZKUpdatesServer() { - ModelZKServer newServer = mock(ModelZKServer.class); - assertDoesNotThrow(() -> listener.setModelZK(newServer)); - } - - @Test - void testProcessReplacesEndpointWithStatusAndCallsUpdateNodeData() { - listener.process(TEST_PATH, mockZk); - - String expectedStatusNode = TEST_PATH.replace("endpoint", "status"); - verify(mockServer).updateNodeData(expectedStatusNode, "INIT", true); - } - - @Test - void testProcessStatusNodePathCorrectness() { - String endpointPath = "/cluster/models/myEngine/endpoint"; - ModelUnavailableListener unavailListener = new ModelUnavailableListener(endpointPath, mockServer); - - unavailListener.process(endpointPath, mockZk); - - verify(mockServer).updateNodeData("/cluster/models/myEngine/status", "INIT", true); - } - - @Test - void testProcessCallsUpdateNodeDataWithCorrectArguments() { - listener.process(TEST_PATH, mockZk); - - String expectedStatusNode = TEST_PATH.replace("endpoint", "status"); - verify(mockServer).updateNodeData( - eq(expectedStatusNode), - eq("INIT"), - eq(true) - ); - } - - @Test - void testProcessAfterSetModelZKUsesNewServer() { - ModelZKServer newServer = mock(ModelZKServer.class); - listener.setModelZK(newServer); - - listener.process(TEST_PATH, mockZk); - - String expectedStatusNode = TEST_PATH.replace("endpoint", "status"); - verify(newServer).updateNodeData(expectedStatusNode, "INIT", true); - verify(mockServer, never()).updateNodeData(anyString(), anyString(), anyBoolean()); - } - - @Test - void testProcessDoesNotInteractWithZooKeeper() { - listener.process(TEST_PATH, mockZk); - - verifyNoInteractions(mockZk); - } -} diff --git a/test/prerna/cluster/util/ModelZKServerUnitTests.java b/test/prerna/cluster/util/ModelZKServerUnitTests.java deleted file mode 100644 index 705c2d309c8..00000000000 --- a/test/prerna/cluster/util/ModelZKServerUnitTests.java +++ /dev/null @@ -1,379 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.*; -import static org.mockito.Mockito.*; - -import java.lang.reflect.Constructor; -import java.lang.reflect.Field; -import java.util.Arrays; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -import org.apache.zookeeper.WatchedEvent; -import org.apache.zookeeper.Watcher.Event.EventType; -import org.apache.zookeeper.Watcher.Event.KeeperState; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Nested; -import org.apache.zookeeper.KeeperException; -import org.apache.zookeeper.ZooKeeper; -import org.apache.zookeeper.data.Stat; -import org.apache.zookeeper.ZooDefs.Ids; -import org.apache.zookeeper.CreateMode; -import org.junit.jupiter.api.Test; -import com.google.gson.Gson; - -class ModelZKServerUnitTests { - - private ModelZKServer server; - - @BeforeEach - void setUp() throws Exception { - Constructor ctor = ModelZKServer.class.getDeclaredConstructor(); - ctor.setAccessible(true); - server = ctor.newInstance(); - } - - @Nested - class ConstantsTests { - @Test void testZkServerConstant() { assertEquals("zk", ModelZKServer.ZK_SERVER); } - @Test void testHostConstant() { assertEquals("host", ModelZKServer.HOST); } - @Test void testTimeoutConstant() { assertEquals("to", ModelZKServer.TIMEOUT); } - @Test void testBootuserConstant() { assertEquals("bu", ModelZKServer.BOOTUSER); } - @Test void testHomeConstant() { assertEquals("zk_home", ModelZKServer.HOME); } - @Test void testAppHomeConstant() { assertEquals("app", ModelZKServer.APP_HOME); } - @Test void testSemossHomeConstant() { assertEquals("sem", ModelZKServer.SEMOSS_HOME); } - @Test void testAvailableConstant() { assertEquals("available", ModelZKServer.AVAILABLE); } - @Test void testModelRootConstant() { assertEquals("/model", ModelZKServer.MODEL_ROOT); } - @Test void testServerPathConstant() { assertEquals("/server", ModelZKServer.SERVER_PATH); } - } - - @Nested - class DefaultFieldValuesTests { - @Test void testDefaultZkServer() { assertEquals("localhost:2181", server.zkServer); } - @Test void testDefaultHost() { assertEquals("localhost:8888", server.host); } - @Test void testDefaultUser() { assertEquals("generic", server.user); } - @Test void testDefaultHome() { assertEquals("/semoss_root", server.home); } - @Test void testDefaultApp() { assertEquals("/app", server.app); } - @Test void testDefaultSemossHome() { assertEquals("/opt/semosshome/", ModelZKServer.semossHome); } - @Test void testDefaultContainer() { assertEquals("/container", server.container); } - @Test void testDefaultId() { assertEquals("RANDOM_ID", server.id); } - @Test void testDefaultCatchupTrue() { assertTrue(server.catchup); } - @Test void testDefaultInitFalse() { assertFalse(server.init); } - @Test void testDefaultPropNull() { assertNull(server.prop); } - @Test void testDefaultZkNull() { assertNull(server.zk); } - @Test void testDefaultEnvNull() { assertNull(server.env); } - @Test void testDefaultListenersEmpty() { assertNotNull(server.listeners); assertTrue(server.listeners.isEmpty()); } - @Test void testDefaultConnectedFalse() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("connected"); f.setAccessible(true); assertFalse((boolean) f.get(server)); - } - @Test void testDefaultGsonNonNull() throws Exception { - Field gsonField = ModelZKServer.class.getDeclaredField("gson"); gsonField.setAccessible(true); - assertNotNull(gsonField.get(server)); assertInstanceOf(Gson.class, gsonField.get(server)); - } - @Test void testDefaultExistingModelsEmpty() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("existingModels"); f.setAccessible(true); - List m = (List) f.get(server); assertNotNull(m); assertTrue(m.isEmpty()); - } - @Test void testDefaultSupportedModelsEmpty() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("supportedModels"); f.setAccessible(true); - List m = (List) f.get(server); assertNotNull(m); assertTrue(m.isEmpty()); - } - @Test void testDefaultModelSMSSEmpty() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("modelSMSS"); f.setAccessible(true); - Map m = (Map) f.get(server); assertNotNull(m); assertTrue(m.isEmpty()); - } - @Test void testDefaultModelLockEmpty() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("modelLock"); f.setAccessible(true); - Map m = (Map) f.get(server); assertNotNull(m); assertTrue(m.isEmpty()); - } - @Test void testDefaultModelClientEmpty() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("modelClient"); f.setAccessible(true); - Map m = (Map) f.get(server); assertNotNull(m); assertTrue(m.isEmpty()); - } - @Test void testDefaultRepeatEmpty() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("repeat"); f.setAccessible(true); - Map m = (Map) f.get(server); assertNotNull(m); assertTrue(m.isEmpty()); - } - @Test void testDefaultClientNull() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("client"); f.setAccessible(true); assertNull(f.get(server)); - } - @Test void testDefaultVersionZero() throws Exception { - Field f = ModelZKServer.class.getDeclaredField("version"); f.setAccessible(true); assertEquals(0, f.getInt(server)); - } - @Test void testDefaultZkClientStaticFieldType() { - try { Field f = ModelZKServer.class.getDeclaredField("zkClient"); assertEquals(ModelZKServer.class, f.getType()); } - catch (NoSuchFieldException e) { fail("zkClient field should exist"); } - } - } - - @Nested - class CanAccomodateTests { - @Test void testBothMapsEmpty_returnsTrue() { assertTrue(server.canAccomodate(new HashMap<>(), new HashMap<>())); } - @Test void testNeedSatisfied_returnsTrue() { - Map cap = new HashMap<>(); Map sub = new HashMap<>(); - sub.put("available", 8.0); cap.put("gpu", sub); - assertTrue(server.canAccomodate(cap, Map.of("gpu", 4))); - } - @Test void testNeedNotSatisfied_returnsFalse() { - Map cap = new HashMap<>(); Map sub = new HashMap<>(); - sub.put("available", 5.0); cap.put("gpu", sub); - assertFalse(server.canAccomodate(cap, Map.of("gpu", 10))); - } - @Test void testCapValueIsNotAMap_returnsFalse() { - Map cap = new HashMap<>(); cap.put("gpu", "notAMap"); - assertFalse(server.canAccomodate(cap, Map.of("gpu", 4))); - } - @Test void testNeedKeyMissingFromCap_returnsFalse() { - assertFalse(server.canAccomodate(new HashMap<>(), Map.of("gpu", 4))); - } - @Test void testNeedExactlyEqualsCapability_returnsTrue() { - Map cap = new HashMap<>(); Map sub = new HashMap<>(); - sub.put("available", 4.0); cap.put("gpu", sub); - assertTrue(server.canAccomodate(cap, Map.of("gpu", 4))); - } - @Test void testMultipleNeedsAllSatisfied_returnsTrue() { - Map cap = new HashMap<>(); - Map gs = new HashMap<>(); gs.put("available", 8.0); cap.put("gpu", gs); - Map ms = new HashMap<>(); ms.put("available", 32.0); cap.put("memory", ms); - Map need = new HashMap<>(); need.put("gpu", 4); need.put("memory", 16); - assertTrue(server.canAccomodate(cap, need)); - } - @Test void testMultipleNeeds_oneFails_returnsFalse() { - Map cap = new HashMap<>(); - Map gs = new HashMap<>(); gs.put("available", 8.0); cap.put("gpu", gs); - Map ms = new HashMap<>(); ms.put("available", 10.0); cap.put("memory", ms); - Map need = new HashMap<>(); need.put("gpu", 4); need.put("memory", 16); - assertFalse(server.canAccomodate(cap, need)); - } - } - - @Nested - class ListenerTests { - @Test void testAddZKListener_addsToListenersMap() { - IModelZKListener m = mock(IModelZKListener.class); when(m.getPath()).thenReturn("/model/test"); - server.addZKListener(m); - assertTrue(server.listeners.containsKey("/model/test")); - assertEquals(1, server.listeners.get("/model/test").size()); - assertSame(m, server.listeners.get("/model/test").get(0)); - } - @Test void testAddZKListener_multipleListenersSamePath() { - IModelZKListener l1 = mock(IModelZKListener.class); when(l1.getPath()).thenReturn("/model/test"); - IModelZKListener l2 = mock(IModelZKListener.class); when(l2.getPath()).thenReturn("/model/test"); - server.addZKListener(l1); server.addZKListener(l2); - assertEquals(2, server.listeners.get("/model/test").size()); - } - @Test void testAddZKListener_differentPaths() { - IModelZKListener l1 = mock(IModelZKListener.class); when(l1.getPath()).thenReturn("/model/t1"); - IModelZKListener l2 = mock(IModelZKListener.class); when(l2.getPath()).thenReturn("/model/t2"); - server.addZKListener(l1); server.addZKListener(l2); - assertTrue(server.listeners.containsKey("/model/t1")); - assertTrue(server.listeners.containsKey("/model/t2")); - } - @Test void testProcessEvent_matching() { - IModelZKListener m = mock(IModelZKListener.class); when(m.getPath()).thenReturn("/model/test"); - when(m.getEvents()).thenReturn(Arrays.asList(EventType.NodeDataChanged)); - server.addZKListener(m); - server.processEvent("/model/test", new WatchedEvent(EventType.NodeDataChanged, KeeperState.SyncConnected, "/model/test")); - verify(m).process("/model/test", null); - } - @Test void testProcessEvent_nonMatchingPath() { - IModelZKListener m = mock(IModelZKListener.class); when(m.getPath()).thenReturn("/model/test"); - when(m.getEvents()).thenReturn(Arrays.asList(EventType.NodeDataChanged)); - server.addZKListener(m); - server.processEvent("/model/other", new WatchedEvent(EventType.NodeDataChanged, KeeperState.SyncConnected, "/model/other")); - verify(m, never()).process(anyString(), any()); - } - @Test void testProcessEvent_wrongEvent() { - IModelZKListener m = mock(IModelZKListener.class); when(m.getPath()).thenReturn("/model/test"); - when(m.getEvents()).thenReturn(Arrays.asList(EventType.NodeCreated)); - server.addZKListener(m); - server.processEvent("/model/test", new WatchedEvent(EventType.NodeDataChanged, KeeperState.SyncConnected, "/model/test")); - verify(m, never()).process(anyString(), any()); - } - @Test void testProcessEvent_noListeners() { - assertDoesNotThrow(() -> server.processEvent("/model/unknown", - new WatchedEvent(EventType.NodeDataChanged, KeeperState.SyncConnected, "/model/unknown"))); - } - @Test void testProcessEvent_multipleListeners() { - IModelZKListener l1 = mock(IModelZKListener.class); when(l1.getPath()).thenReturn("/model/test"); - IModelZKListener l2 = mock(IModelZKListener.class); when(l2.getPath()).thenReturn("/model/test"); - when(l1.getEvents()).thenReturn(Arrays.asList(EventType.NodeDataChanged)); - when(l2.getEvents()).thenReturn(Arrays.asList(EventType.NodeDataChanged)); - server.addZKListener(l1); server.addZKListener(l2); - server.processEvent("/model/test", new WatchedEvent(EventType.NodeDataChanged, KeeperState.SyncConnected, "/model/test")); - verify(l1).process("/model/test", null); verify(l2).process("/model/test", null); - } - @Test void testProcessEvent_multipleEventTypes() { - IModelZKListener m = mock(IModelZKListener.class); when(m.getPath()).thenReturn("/model/test"); - when(m.getEvents()).thenReturn(Arrays.asList(EventType.NodeCreated, EventType.NodeDeleted)); - server.addZKListener(m); - server.processEvent("/model/test", new WatchedEvent(EventType.NodeDeleted, KeeperState.SyncConnected, "/model/test")); - verify(m).process("/model/test", null); - } - } - - @Nested - class MockZKTests { - private ZooKeeper mockZk; - @BeforeEach void setUpMockZk() { mockZk = mock(ZooKeeper.class); server.zk = mockZk; } - @Test void testGetNodeData_returnsStringFromBytes() throws Exception { - when(mockZk.getData(eq("/model/test"), eq(true), any(Stat.class))).thenReturn("hello world".getBytes("UTF-8")); - assertEquals("hello world", server.getNodeData("/model/test")); - verify(mockZk).getData(eq("/model/test"), eq(true), any(Stat.class)); - } - @Test void testGetNodeData_exceptionReturnsNull() throws Exception { - when(mockZk.getData(eq("/model/test"), eq(true), any(Stat.class))).thenThrow(new KeeperException.NoNodeException()); - assertNull(server.getNodeData("/model/test")); - } - @Test void testGetNodeData_emptyBytes() throws Exception { - when(mockZk.getData(eq("/model/empty"), eq(true), any(Stat.class))).thenReturn(new byte[0]); - assertEquals("", server.getNodeData("/model/empty")); - } - @Test void testUpdateNodeData_createTrue_nodeDoesNotExist() throws Exception { - when(mockZk.exists(eq("/model/test"), eq(false))).thenReturn(null); - server.updateNodeData("/model/test", "data", true); - verify(mockZk).create(eq("/model/test"), eq("data".getBytes()), eq(Ids.OPEN_ACL_UNSAFE), eq(CreateMode.EPHEMERAL), any(Stat.class)); - verify(mockZk, never()).setData(anyString(), any(byte[].class), anyInt()); - } - @Test void testUpdateNodeData_createFalse_doesNothing() throws Exception { - when(mockZk.exists(eq("/model/test"), eq(false))).thenReturn(null); - server.updateNodeData("/model/test", "data", false); - verify(mockZk, never()).create(anyString(), any(byte[].class), anyList(), any(CreateMode.class), any(Stat.class)); - verify(mockZk, never()).setData(anyString(), any(byte[].class), anyInt()); - } - @Test void testUpdateNodeData_nodeExists_callsSetData() throws Exception { - when(mockZk.exists(eq("/model/test"), eq(false))).thenReturn(new Stat()); - server.updateNodeData("/model/test", "updated", true); - verify(mockZk).setData(eq("/model/test"), eq("updated".getBytes()), eq(-1)); - } - @Test void testUpdateNodeData_keeperException() throws Exception { - when(mockZk.exists(eq("/model/test"), eq(false))).thenThrow(new KeeperException.ConnectionLossException()); - assertDoesNotThrow(() -> server.updateNodeData("/model/test", "data", true)); - } - @Test void testGetChildren_watchTrue() throws Exception { - List expected = Arrays.asList("child1", "child2"); - when(mockZk.getChildren(eq("/model"), eq(true))).thenReturn(expected); - assertEquals(expected, server.getChildren("/model", true)); - } - @Test void testGetChildren_watchFalse() throws Exception { - when(mockZk.getChildren(eq("/server"), eq(false))).thenReturn(Arrays.asList("a", "b", "c")); - assertEquals(Arrays.asList("a", "b", "c"), server.getChildren("/server", false)); - } - @Test void testGetChildren_keeperException() throws Exception { - when(mockZk.getChildren(eq("/model"), eq(true))).thenThrow(new KeeperException.ConnectionLossException()); - assertThrows(IllegalStateException.class, () -> server.getChildren("/model", true)); - } - @Test void testGetChildren_interruptedException() throws Exception { - when(mockZk.getChildren(eq("/model"), eq(true))).thenThrow(new InterruptedException("test")); - assertThrows(IllegalStateException.class, () -> server.getChildren("/model", true)); - } - @Test void testGetChildren_emptyList() throws Exception { - when(mockZk.getChildren(eq("/model"), eq(false))).thenReturn(Arrays.asList()); - assertTrue(server.getChildren("/model", false).isEmpty()); - } - @Test void testAddServer_pathDoesNotExist() throws Exception { - when(mockZk.exists(eq(ModelZKServer.SERVER_PATH), eq(false))).thenReturn(null); - when(mockZk.create(anyString(), any(byte[].class), anyList(), any(CreateMode.class))).thenReturn(""); - server.addServer(); - verify(mockZk).create(eq(ModelZKServer.SERVER_PATH), eq("1".getBytes()), eq(Ids.OPEN_ACL_UNSAFE), eq(CreateMode.PERSISTENT)); - verify(mockZk).create(eq(ModelZKServer.SERVER_PATH + "/" + server.id), eq(server.id.getBytes()), eq(Ids.OPEN_ACL_UNSAFE), eq(CreateMode.EPHEMERAL)); - } - @Test void testAddServer_pathExists() throws Exception { - when(mockZk.exists(eq(ModelZKServer.SERVER_PATH), eq(false))).thenReturn(new Stat()); - when(mockZk.create(anyString(), any(byte[].class), anyList(), any(CreateMode.class))).thenReturn(""); - server.addServer(); - verify(mockZk, never()).create(eq(ModelZKServer.SERVER_PATH), any(byte[].class), anyList(), eq(CreateMode.PERSISTENT)); - verify(mockZk).create(eq(ModelZKServer.SERVER_PATH + "/" + server.id), eq(server.id.getBytes()), eq(Ids.OPEN_ACL_UNSAFE), eq(CreateMode.EPHEMERAL)); - } - @Test void testAddServer_keeperException() throws Exception { - when(mockZk.exists(eq(ModelZKServer.SERVER_PATH), eq(false))).thenThrow(new KeeperException.ConnectionLossException()); - assertDoesNotThrow(() -> server.addServer()); - } - @Test void testProcess_delegatesToProcessEvent() { - IModelZKListener m = mock(IModelZKListener.class); when(m.getPath()).thenReturn("/model/test"); - when(m.getEvents()).thenReturn(Arrays.asList(EventType.NodeDataChanged)); - server.addZKListener(m); - server.process(new WatchedEvent(EventType.NodeDataChanged, KeeperState.SyncConnected, "/model/test")); - verify(m).process("/model/test", mockZk); - } - @Test void testProcess_nullPath() { - assertDoesNotThrow(() -> server.process(new WatchedEvent(EventType.None, KeeperState.SyncConnected, null))); - } - @Test void testProcess_nodeCreated() { - IModelZKListener m = mock(IModelZKListener.class); when(m.getPath()).thenReturn("/model/new"); - when(m.getEvents()).thenReturn(Arrays.asList(EventType.NodeCreated)); - server.addZKListener(m); - server.process(new WatchedEvent(EventType.NodeCreated, KeeperState.SyncConnected, "/model/new")); - verify(m).process("/model/new", mockZk); - } - @Test void testProcess_nodeDeleted() { - IModelZKListener m = mock(IModelZKListener.class); when(m.getPath()).thenReturn("/model/gone"); - when(m.getEvents()).thenReturn(Arrays.asList(EventType.NodeDeleted)); - server.addZKListener(m); - server.process(new WatchedEvent(EventType.NodeDeleted, KeeperState.SyncConnected, "/model/gone")); - verify(m).process("/model/gone", mockZk); - } - @Test void testWatch4Data_callsGetData() throws Exception { - when(mockZk.getData(eq("/model/test"), eq(true), any(Stat.class))).thenReturn("data".getBytes("UTF-8")); - server.watch4Data("/model/test"); - verify(mockZk).getData(eq("/model/test"), eq(true), any(Stat.class)); - } - @Test void testWatch4Data_keeperException() throws Exception { - when(mockZk.getData(eq("/model/m"), eq(true), any(Stat.class))).thenThrow(new KeeperException.NoNodeException()); - assertDoesNotThrow(() -> server.watch4Data("/model/m")); - } - @Test void testWatch4Data_interruptedException() throws Exception { - when(mockZk.getData(eq("/model/t"), eq(true), any(Stat.class))).thenThrow(new InterruptedException("t")); - assertDoesNotThrow(() -> server.watch4Data("/model/t")); - } - @Test void testWatch4Children_callsGetChildren() throws Exception { - when(mockZk.getChildren(eq("/model"), eq(true))).thenReturn(Arrays.asList("c1")); - server.watch4Children("/model"); - verify(mockZk).getChildren(eq("/model"), eq(true)); - } - @Test void testWatch4Children_keeperException() throws Exception { - when(mockZk.getChildren(eq("/model/m"), eq(true))).thenThrow(new KeeperException.NoNodeException()); - assertDoesNotThrow(() -> server.watch4Children("/model/m")); - } - @Test void testWatch4Children_interruptedException() throws Exception { - when(mockZk.getChildren(eq("/model/t"), eq(true))).thenThrow(new InterruptedException("t")); - assertDoesNotThrow(() -> server.watch4Children("/model/t")); - } - @Test void testReconnect_whenZkClientIsNull() throws Exception { - Field zkClientField = ModelZKServer.class.getDeclaredField("zkClient"); - zkClientField.setAccessible(true); - Object original = zkClientField.get(null); - try { zkClientField.set(null, null); assertThrows(NullPointerException.class, () -> server.reconnect()); } - finally { zkClientField.set(null, original); } - } - } -} diff --git a/test/prerna/cluster/util/RemoteClientServerZKRESTProxyUnitTests.java b/test/prerna/cluster/util/RemoteClientServerZKRESTProxyUnitTests.java deleted file mode 100644 index a7bc2f9161f..00000000000 --- a/test/prerna/cluster/util/RemoteClientServerZKRESTProxyUnitTests.java +++ /dev/null @@ -1,1304 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.*; -import static org.mockito.Mockito.*; - -import java.io.IOException; -import java.lang.reflect.Constructor; -import java.lang.reflect.Field; -import java.lang.reflect.Method; -import java.util.List; -import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.ScheduledExecutorService; - -import org.apache.http.HttpEntity; -import org.apache.http.StatusLine; -import org.apache.http.client.methods.CloseableHttpResponse; -import org.apache.http.client.methods.HttpGet; -import org.apache.http.client.methods.HttpHead; -import org.apache.http.client.methods.HttpUriRequest; -import org.apache.http.entity.StringEntity; -import org.apache.http.impl.client.CloseableHttpClient; -import org.apache.http.impl.client.HttpClientBuilder; -import org.apache.http.impl.client.HttpClients; -import org.json.JSONObject; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.MockedStatic; - -import prerna.engine.api.RemoteModelStateEnum; - -class RemoteClientServerZKRESTProxyUnitTests { - - private RemoteClientServerZKRESTProxy proxy; - - @BeforeEach - void setUp() throws Exception { - Constructor ctor = - RemoteClientServerZKRESTProxy.class.getDeclaredConstructor(); - ctor.setAccessible(true); - proxy = ctor.newInstance(); - - setField("zkRestProxyBaseUrl", "http://localhost:8080/"); - } - - @AfterEach - void tearDown() throws Exception { - Field instanceField = RemoteClientServerZKRESTProxy.class.getDeclaredField("instance"); - instanceField.setAccessible(true); - instanceField.set(null, null); - } - - // ========== Reflection helpers ========== - - private void setField(String fieldName, Object value) throws Exception { - Field f = RemoteClientServerZKRESTProxy.class.getDeclaredField(fieldName); - f.setAccessible(true); - f.set(proxy, value); - } - - @SuppressWarnings("unchecked") - private ConcurrentMap getModelStates() throws Exception { - Field f = RemoteClientServerZKRESTProxy.class.getDeclaredField("modelStates"); - f.setAccessible(true); - return (ConcurrentMap) f.get(proxy); - } - - @SuppressWarnings("unchecked") - private ConcurrentMap getModelClusterIps() throws Exception { - Field f = RemoteClientServerZKRESTProxy.class.getDeclaredField("modelClusterIps"); - f.setAccessible(true); - return (ConcurrentMap) f.get(proxy); - } - - @SuppressWarnings("unchecked") - private ConcurrentMap getModelNames() throws Exception { - Field f = RemoteClientServerZKRESTProxy.class.getDeclaredField("modelNames"); - f.setAccessible(true); - return (ConcurrentMap) f.get(proxy); - } - - private ScheduledExecutorService getScheduler() throws Exception { - Field f = RemoteClientServerZKRESTProxy.class.getDeclaredField("scheduler"); - f.setAccessible(true); - return (ScheduledExecutorService) f.get(proxy); - } - - // ========== HTTP mock helpers ========== - - private CloseableHttpResponse createMockResponse(int statusCode, String body) throws Exception { - CloseableHttpResponse response = mock(CloseableHttpResponse.class); - StatusLine statusLine = mock(StatusLine.class); - when(statusLine.getStatusCode()).thenReturn(statusCode); - when(response.getStatusLine()).thenReturn(statusLine); - - if (body != null) { - when(response.getEntity()).thenReturn(new StringEntity(body)); - } else { - when(response.getEntity()).thenReturn(null); - } - - return response; - } - - private MockedStatic setupHttpMock(CloseableHttpClient mockClient) { - HttpClientBuilder mockBuilder = mock(HttpClientBuilder.class); - when(mockBuilder.setDefaultRequestConfig(any())).thenReturn(mockBuilder); - when(mockBuilder.build()).thenReturn(mockClient); - - MockedStatic mockedHttpClients = mockStatic(HttpClients.class); - mockedHttpClients.when(HttpClients::custom).thenReturn(mockBuilder); - return mockedHttpClients; - } - - // ========== Private method invocation helpers ========== - - private Object invokePrivate(String methodName, Class[] paramTypes, Object... args) throws Exception { - Method m = RemoteClientServerZKRESTProxy.class.getDeclaredMethod(methodName, paramTypes); - m.setAccessible(true); - return m.invoke(proxy, args); - } - - // ========== Tests: getModelState ========== - - @Test - void testGetModelState_activeModel() throws Exception { - getModelStates().put("model-1", RemoteModelStateEnum.ACTIVE); - assertEquals(RemoteModelStateEnum.ACTIVE, proxy.getModelState("model-1")); - } - - @Test - void testGetModelState_warmingModel() throws Exception { - getModelStates().put("model-2", RemoteModelStateEnum.WARMING); - assertEquals(RemoteModelStateEnum.WARMING, proxy.getModelState("model-2")); - } - - @Test - void testGetModelState_coldModel() throws Exception { - getModelStates().put("model-3", RemoteModelStateEnum.COLD); - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("model-3")); - } - - @Test - void testGetModelState_failedModel() throws Exception { - getModelStates().put("model-4", RemoteModelStateEnum.FAILED); - assertEquals(RemoteModelStateEnum.FAILED, proxy.getModelState("model-4")); - } - - @Test - void testGetModelState_unknownModelDefaultsToCold() { - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("nonexistent")); - } - - @Test - void testGetModelState_nullModelId_throwsNPE() { - assertThrows(NullPointerException.class, () -> proxy.getModelState(null)); - } - - @Test - void testGetModelState_emptyModelId() { - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("")); - } - - // ========== Tests: getModelName ========== - - @Test - void testGetModelName_known() throws Exception { - getModelNames().put("model-1", "GPT-4"); - assertEquals("GPT-4", proxy.getModelName("model-1")); - } - - @Test - void testGetModelName_unknown() { - assertNull(proxy.getModelName("nonexistent")); - } - - @Test - void testGetModelName_emptyString() throws Exception { - getModelNames().put("model-1", ""); - assertEquals("", proxy.getModelName("model-1")); - } - - @Test - void testGetModelName_multipleModels() throws Exception { - getModelNames().put("m1", "Model-A"); - getModelNames().put("m2", "Model-B"); - getModelNames().put("m3", "Model-C"); - assertEquals("Model-A", proxy.getModelName("m1")); - assertEquals("Model-B", proxy.getModelName("m2")); - assertEquals("Model-C", proxy.getModelName("m3")); - } - - // ========== Tests: getModelClusterIp (cache hit) ========== - - @Test - void testGetModelClusterIp_cached() throws Exception { - getModelClusterIps().put("model-1", "10.0.0.1"); - assertEquals("10.0.0.1", proxy.getModelClusterIp("model-1")); - } - - @Test - void testGetModelClusterIp_multipleCached() throws Exception { - getModelClusterIps().put("m1", "10.0.0.1"); - getModelClusterIps().put("m2", "10.0.0.2"); - assertEquals("10.0.0.1", proxy.getModelClusterIp("m1")); - assertEquals("10.0.0.2", proxy.getModelClusterIp("m2")); - } - - @Test - void testGetModelClusterIp_cacheMiss_httpReturnsData() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject data = new JSONObject(); - data.put("ip", "10.0.0.5"); - data.put("model_name", "TestModel"); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", data); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String ip = proxy.getModelClusterIp("model-new"); - assertEquals("10.0.0.5", ip); - // Verify it also cached the results - assertEquals("10.0.0.5", getModelClusterIps().get("model-new")); - assertEquals("TestModel", getModelNames().get("model-new")); - } - } - - @Test - void testGetModelClusterIp_cacheMiss_httpReturns404() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String ip = proxy.getModelClusterIp("missing-model"); - assertNull(ip); - } - } - - @Test - void testGetModelClusterIp_cacheMiss_httpReturns500() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(500, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String ip = proxy.getModelClusterIp("error-model"); - assertNull(ip); - } - } - - @Test - void testGetModelClusterIp_cacheMiss_httpThrowsException() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpGet.class))).thenThrow(new IOException("Connection refused")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String ip = proxy.getModelClusterIp("bad-model"); - assertNull(ip); - } - } - - // ========== Tests: getModelScalerIp ========== - - @Test - void testGetModelScalerIp_doesNotThrow() { - assertDoesNotThrow(() -> proxy.getModelScalerIp()); - } - - // ========== Tests: close ========== - - @Test - void testClose_shutsDownScheduler() throws Exception { - proxy.close(); - assertTrue(getScheduler().isShutdown()); - } - - @Test - void testClose_calledTwice_noException() { - proxy.close(); - assertDoesNotThrow(() -> proxy.close()); - } - - @Test - void testClose_schedulerIsShutdownAfterClose() throws Exception { - assertFalse(getScheduler().isShutdown()); - proxy.close(); - assertTrue(getScheduler().isShutdown()); - } - - // ========== Tests: isModelWarming (via checkZNodeExists) ========== - - @Test - void testIsModelWarming_exists() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertTrue(proxy.isModelWarming("model-1")); - } - } - - @Test - void testIsModelWarming_notExists() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertFalse(proxy.isModelWarming("model-1")); - } - } - - @Test - void testIsModelWarming_serverError() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(500, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertFalse(proxy.isModelWarming("model-1")); - } - } - - @Test - void testIsModelWarming_ioException() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpHead.class))).thenThrow(new IOException("Timeout")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertFalse(proxy.isModelWarming("model-1")); - } - } - - // ========== Tests: isModelActive (via checkZNodeExists) ========== - - @Test - void testIsModelActive_exists() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertTrue(proxy.isModelActive("model-1")); - } - } - - @Test - void testIsModelActive_notExists() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertFalse(proxy.isModelActive("model-1")); - } - } - - @Test - void testIsModelActive_ioException() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpHead.class))).thenThrow(new IOException("Refused")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertFalse(proxy.isModelActive("model-1")); - } - } - - // ========== Tests: getZNodeDataAsJson (private) ========== - - @Test - void testGetZNodeDataAsJson_success() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject data = new JSONObject(); - data.put("ip", "10.0.0.1"); - data.put("model_name", "TestModel"); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", data); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "/models/active/m1"); - assertNotNull(result); - assertEquals("10.0.0.1", result.getString("ip")); - assertEquals("TestModel", result.getString("model_name")); - } - } - - @Test - void testGetZNodeDataAsJson_404() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "/models/active/m1"); - assertNull(result); - } - } - - @Test - void testGetZNodeDataAsJson_500() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(500, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeDataAsJson_nullDataField() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", JSONObject.NULL); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeDataAsJson_noDataField() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject responseBody = new JSONObject(); - responseBody.put("other", "value"); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeDataAsJson_nullEntity() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeDataAsJson_pathWithLeadingSlash() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject data = new JSONObject(); - data.put("ip", "1.2.3.4"); - data.put("model_name", "Slash"); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", data); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "/models/active/m1"); - assertNotNull(result); - assertEquals("1.2.3.4", result.getString("ip")); - } - } - - @Test - void testGetZNodeDataAsJson_pathWithoutLeadingSlash() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject data = new JSONObject(); - data.put("ip", "5.6.7.8"); - data.put("model_name", "NoSlash"); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", data); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "models/active/m1"); - assertNotNull(result); - assertEquals("5.6.7.8", result.getString("ip")); - } - } - - @Test - void testGetZNodeDataAsJson_ioException() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpGet.class))).thenThrow(new IOException("Network error")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - JSONObject result = (JSONObject) invokePrivate("getZNodeDataAsJson", - new Class[]{String.class}, "/models/active/m1"); - assertNull(result); - } - } - - // ========== Tests: getZNodeData (private) ========== - - @Test - void testGetZNodeData_success() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", "some-string-data"); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String result = (String) invokePrivate("getZNodeData", - new Class[]{String.class}, "/services/kube-model-deployer"); - assertEquals("some-string-data", result); - } - } - - @Test - void testGetZNodeData_404() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String result = (String) invokePrivate("getZNodeData", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeData_500() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(500, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String result = (String) invokePrivate("getZNodeData", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeData_noDataField() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject responseBody = new JSONObject(); - responseBody.put("other", "value"); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String result = (String) invokePrivate("getZNodeData", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeData_nullDataField() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", JSONObject.NULL); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String result = (String) invokePrivate("getZNodeData", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeData_nullEntity() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String result = (String) invokePrivate("getZNodeData", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - @Test - void testGetZNodeData_ioException() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpGet.class))).thenThrow(new IOException("Network error")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - String result = (String) invokePrivate("getZNodeData", - new Class[]{String.class}, "/some/path"); - assertNull(result); - } - } - - // ========== Tests: getZNodeChildren (private) ========== - - @Test - void testGetZNodeChildren_success() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject responseBody = new JSONObject(); - responseBody.put("children", new org.json.JSONArray(List.of("child1", "child2", "child3"))); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - @SuppressWarnings("unchecked") - List result = (List) invokePrivate("getZNodeChildren", - new Class[]{String.class}, "/models/active"); - assertEquals(3, result.size()); - assertTrue(result.contains("child1")); - assertTrue(result.contains("child2")); - assertTrue(result.contains("child3")); - } - } - - @Test - void testGetZNodeChildren_emptyChildren() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject responseBody = new JSONObject(); - responseBody.put("children", new org.json.JSONArray()); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - @SuppressWarnings("unchecked") - List result = (List) invokePrivate("getZNodeChildren", - new Class[]{String.class}, "/models/active"); - assertTrue(result.isEmpty()); - } - } - - @Test - void testGetZNodeChildren_noChildrenField() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", "something"); - - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - @SuppressWarnings("unchecked") - List result = (List) invokePrivate("getZNodeChildren", - new Class[]{String.class}, "/models/active"); - assertTrue(result.isEmpty()); - } - } - - @Test - void testGetZNodeChildren_404() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - @SuppressWarnings("unchecked") - List result = (List) invokePrivate("getZNodeChildren", - new Class[]{String.class}, "/models/active"); - assertTrue(result.isEmpty()); - } - } - - @Test - void testGetZNodeChildren_500() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(500, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - @SuppressWarnings("unchecked") - List result = (List) invokePrivate("getZNodeChildren", - new Class[]{String.class}, "/models/active"); - assertTrue(result.isEmpty()); - } - } - - @Test - void testGetZNodeChildren_ioException() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpGet.class))).thenThrow(new IOException("Connection refused")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - @SuppressWarnings("unchecked") - List result = (List) invokePrivate("getZNodeChildren", - new Class[]{String.class}, "/models/active"); - assertTrue(result.isEmpty()); - } - } - - // ========== Tests: checkZNodeExists (private) ========== - - @Test - void testCheckZNodeExists_true() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - boolean result = (boolean) invokePrivate("checkZNodeExists", - new Class[]{String.class}, "/models/warming/m1"); - assertTrue(result); - } - } - - @Test - void testCheckZNodeExists_false_404() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - boolean result = (boolean) invokePrivate("checkZNodeExists", - new Class[]{String.class}, "/models/warming/m1"); - assertFalse(result); - } - } - - @Test - void testCheckZNodeExists_false_serverError() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(500, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - boolean result = (boolean) invokePrivate("checkZNodeExists", - new Class[]{String.class}, "/models/warming/m1"); - assertFalse(result); - } - } - - @Test - void testCheckZNodeExists_false_ioException() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpHead.class))).thenThrow(new IOException("Timeout")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - boolean result = (boolean) invokePrivate("checkZNodeExists", - new Class[]{String.class}, "/models/warming/m1"); - assertFalse(result); - } - } - - @Test - void testCheckZNodeExists_pathStripsLeadingSlash() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - boolean result = (boolean) invokePrivate("checkZNodeExists", - new Class[]{String.class}, "/test/path"); - assertTrue(result); - } - } - - // ========== Tests: getActiveModels ========== - - @Test - void testGetActiveModels_withCachedNames() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - // getZNodeChildren returns two model IDs - JSONObject childrenResponse = new JSONObject(); - childrenResponse.put("children", new org.json.JSONArray(List.of("m1", "m2"))); - CloseableHttpResponse childrenResp = createMockResponse(200, childrenResponse.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(childrenResp); - - // Pre-populate caches - getModelNames().put("m1", "Model-A"); - getModelNames().put("m2", "Model-B"); - getModelStates().put("m1", RemoteModelStateEnum.ACTIVE); - getModelStates().put("m2", RemoteModelStateEnum.ACTIVE); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - List result = proxy.getActiveModels(); - assertEquals(2, result.size()); - assertEquals("m1", result.get(0).getId()); - assertEquals("Model-A", result.get(0).getName()); - assertEquals(RemoteModelStateEnum.ACTIVE, result.get(0).getState()); - assertEquals("m2", result.get(1).getId()); - assertEquals("Model-B", result.get(1).getName()); - } - } - - @Test - void testGetActiveModels_empty() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject childrenResponse = new JSONObject(); - childrenResponse.put("children", new org.json.JSONArray()); - CloseableHttpResponse childrenResp = createMockResponse(200, childrenResponse.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(childrenResp); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - List result = proxy.getActiveModels(); - assertTrue(result.isEmpty()); - } - } - - @Test - void testGetActiveModels_unknownNameFallsBackToHttp() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - // First call: getZNodeChildren - JSONObject childrenResponse = new JSONObject(); - childrenResponse.put("children", new org.json.JSONArray(List.of("m1"))); - CloseableHttpResponse childrenResp = createMockResponse(200, childrenResponse.toString()); - - // Second call: getZNodeDataAsJson for m1 - JSONObject data = new JSONObject(); - data.put("ip", "10.0.0.1"); - data.put("model_name", "FetchedModel"); - JSONObject dataResponse = new JSONObject(); - dataResponse.put("data", data); - CloseableHttpResponse dataResp = createMockResponse(200, dataResponse.toString()); - - when(mockClient.execute(any(HttpGet.class))) - .thenReturn(childrenResp) - .thenReturn(dataResp); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - List result = proxy.getActiveModels(); - assertEquals(1, result.size()); - assertEquals("m1", result.get(0).getId()); - assertEquals("FetchedModel", result.get(0).getName()); - // Verify it cached the name - assertEquals("FetchedModel", getModelNames().get("m1")); - } - } - - @Test - void testGetActiveModels_unknownNameAndHttpFails() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject childrenResponse = new JSONObject(); - childrenResponse.put("children", new org.json.JSONArray(List.of("m1"))); - CloseableHttpResponse childrenResp = createMockResponse(200, childrenResponse.toString()); - - CloseableHttpResponse dataResp = createMockResponse(404, null); - - when(mockClient.execute(any(HttpGet.class))) - .thenReturn(childrenResp) - .thenReturn(dataResp); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - List result = proxy.getActiveModels(); - assertEquals(1, result.size()); - assertEquals("m1", result.get(0).getId()); - assertEquals("Unknown", result.get(0).getName()); - } - } - - @Test - void testGetActiveModels_httpExceptionReturnsEmptyList() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpGet.class))).thenThrow(new IOException("Network error")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - List result = proxy.getActiveModels(); - assertTrue(result.isEmpty()); - } - } - - // ========== Tests: getWarmingModels ========== - - @Test - void testGetWarmingModels_withCachedNames() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject childrenResponse = new JSONObject(); - childrenResponse.put("children", new org.json.JSONArray(List.of("w1", "w2"))); - CloseableHttpResponse childrenResp = createMockResponse(200, childrenResponse.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(childrenResp); - - getModelNames().put("w1", "Warming-A"); - getModelStates().put("w1", RemoteModelStateEnum.WARMING); - getModelStates().put("w2", RemoteModelStateEnum.WARMING); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - List result = proxy.getWarmingModels(); - assertEquals(2, result.size()); - assertEquals("w1", result.get(0).getId()); - assertEquals("Warming-A", result.get(0).getName()); - assertEquals(RemoteModelStateEnum.WARMING, result.get(0).getState()); - assertEquals("w2", result.get(1).getId()); - assertEquals("Warming...", result.get(1).getName()); // Unknown name defaults - } - } - - @Test - void testGetWarmingModels_empty() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject childrenResponse = new JSONObject(); - childrenResponse.put("children", new org.json.JSONArray()); - CloseableHttpResponse childrenResp = createMockResponse(200, childrenResponse.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(childrenResp); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - List result = proxy.getWarmingModels(); - assertTrue(result.isEmpty()); - } - } - - @Test - void testGetWarmingModels_httpExceptionReturnsEmptyList() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpGet.class))).thenThrow(new IOException("Network error")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - List result = proxy.getWarmingModels(); - assertTrue(result.isEmpty()); - } - } - - // ========== Tests: waitForModelActive ========== - - @Test - void testWaitForModelActive_alreadyActive() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertTrue(proxy.waitForModelActive("model-1", 5000)); - assertEquals(RemoteModelStateEnum.ACTIVE, getModelStates().get("model-1")); - } - } - - @Test - void testWaitForModelActive_timeout() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - // Use very short timeout to avoid long waits - assertFalse(proxy.waitForModelActive("model-1", 100)); - } - } - - // ========== Tests: waitForState ========== - - @Test - void testWaitForState_alreadyInDesiredState() throws Exception { - getModelStates().put("model-1", RemoteModelStateEnum.ACTIVE); - assertTrue(proxy.waitForState("model-1", RemoteModelStateEnum.ACTIVE, 5000)); - } - - @Test - void testWaitForState_waitingForCold_defaultsCold() { - // Unknown model defaults to COLD - assertTrue(proxy.waitForState("nonexistent", RemoteModelStateEnum.COLD, 5000)); - } - - @Test - void testWaitForState_warmingModelWaitsForActive_timeout() throws Exception { - getModelStates().put("model-1", RemoteModelStateEnum.WARMING); - - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertFalse(proxy.waitForState("model-1", RemoteModelStateEnum.ACTIVE, 100)); - } - } - - @Test - void testWaitForState_coldModelBecomesActive() throws Exception { - // Model is COLD, but isModelActive returns true - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpHead.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertTrue(proxy.waitForState("model-1", RemoteModelStateEnum.ACTIVE, 5000)); - assertEquals(RemoteModelStateEnum.ACTIVE, getModelStates().get("model-1")); - } - } - - // ========== Tests: refreshModelStates (private) ========== - - @Test - void testRefreshModelStates_updatesStates() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - // First call: getZNodeChildren for ACTIVE_PATH returns ["m1"] - JSONObject activeChildren = new JSONObject(); - activeChildren.put("children", new org.json.JSONArray(List.of("m1"))); - CloseableHttpResponse activeChildrenResp = createMockResponse(200, activeChildren.toString()); - - // Second call: getZNodeChildren for WARMING_PATH returns ["m2"] - JSONObject warmingChildren = new JSONObject(); - warmingChildren.put("children", new org.json.JSONArray(List.of("m2"))); - CloseableHttpResponse warmingChildrenResp = createMockResponse(200, warmingChildren.toString()); - - // Third call: getZNodeDataAsJson for m1 - JSONObject data = new JSONObject(); - data.put("ip", "10.0.0.1"); - data.put("model_name", "Model-A"); - JSONObject dataResponse = new JSONObject(); - dataResponse.put("data", data); - CloseableHttpResponse dataResp = createMockResponse(200, dataResponse.toString()); - - when(mockClient.execute(any(HttpGet.class))) - .thenReturn(activeChildrenResp) - .thenReturn(warmingChildrenResp) - .thenReturn(dataResp); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - invokePrivate("refreshModelStates", new Class[]{}); - assertEquals(RemoteModelStateEnum.ACTIVE, getModelStates().get("m1")); - assertEquals(RemoteModelStateEnum.WARMING, getModelStates().get("m2")); - assertEquals("10.0.0.1", getModelClusterIps().get("m1")); - assertEquals("Model-A", getModelNames().get("m1")); - } - } - - @Test - void testRefreshModelStates_staleEntriesBecomeCold() throws Exception { - // Pre-populate with a model that won't appear in refresh - getModelStates().put("stale-model", RemoteModelStateEnum.ACTIVE); - - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject activeChildren = new JSONObject(); - activeChildren.put("children", new org.json.JSONArray()); - CloseableHttpResponse activeResp = createMockResponse(200, activeChildren.toString()); - - JSONObject warmingChildren = new JSONObject(); - warmingChildren.put("children", new org.json.JSONArray()); - CloseableHttpResponse warmingResp = createMockResponse(200, warmingChildren.toString()); - - when(mockClient.execute(any(HttpGet.class))) - .thenReturn(activeResp) - .thenReturn(warmingResp); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - invokePrivate("refreshModelStates", new Class[]{}); - assertEquals(RemoteModelStateEnum.COLD, getModelStates().get("stale-model")); - } - } - - @Test - void testRefreshModelStates_activeOverridesWarming() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - // Model appears in both active and warming - JSONObject activeChildren = new JSONObject(); - activeChildren.put("children", new org.json.JSONArray(List.of("m1"))); - CloseableHttpResponse activeResp = createMockResponse(200, activeChildren.toString()); - - JSONObject warmingChildren = new JSONObject(); - warmingChildren.put("children", new org.json.JSONArray(List.of("m1"))); - CloseableHttpResponse warmingResp = createMockResponse(200, warmingChildren.toString()); - - JSONObject data = new JSONObject(); - data.put("ip", "10.0.0.1"); - data.put("model_name", "Model-A"); - JSONObject dataResponse = new JSONObject(); - dataResponse.put("data", data); - CloseableHttpResponse dataResp = createMockResponse(200, dataResponse.toString()); - - when(mockClient.execute(any(HttpGet.class))) - .thenReturn(activeResp) - .thenReturn(warmingResp) - .thenReturn(dataResp); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - invokePrivate("refreshModelStates", new Class[]{}); - // Should still be ACTIVE because active is processed first - assertEquals(RemoteModelStateEnum.ACTIVE, getModelStates().get("m1")); - } - } - - // ========== Tests: updateModelInfoFromZNode (private) ========== - - @Test - void testUpdateModelInfoFromZNode_success() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - JSONObject data = new JSONObject(); - data.put("ip", "10.0.0.99"); - data.put("model_name", "UpdatedModel"); - JSONObject responseBody = new JSONObject(); - responseBody.put("data", data); - CloseableHttpResponse mockResponse = createMockResponse(200, responseBody.toString()); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - invokePrivate("updateModelInfoFromZNode", - new Class[]{String.class, String.class}, "m1", "/models/active/m1"); - assertEquals("10.0.0.99", getModelClusterIps().get("m1")); - assertEquals("UpdatedModel", getModelNames().get("m1")); - } - } - - @Test - void testUpdateModelInfoFromZNode_nullData() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(404, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - invokePrivate("updateModelInfoFromZNode", - new Class[]{String.class, String.class}, "m1", "/models/active/m1"); - assertNull(getModelClusterIps().get("m1")); - assertNull(getModelNames().get("m1")); - } - } - - // ========== Tests: validateConnection (private) ========== - - @Test - void testValidateConnection_success() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(200, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertDoesNotThrow(() -> invokePrivate("validateConnection", new Class[]{})); - } - } - - @Test - void testValidateConnection_nonOkStatus() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = createMockResponse(503, null); - when(mockClient.execute(any(HttpGet.class))).thenReturn(mockResponse); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - Exception ex = assertThrows(Exception.class, - () -> invokePrivate("validateConnection", new Class[]{})); - assertTrue(ex.getCause().getMessage().contains("health check failed")); - } - } - - @Test - void testValidateConnection_ioException() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpGet.class))).thenThrow(new IOException("Connection refused")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - Exception ex = assertThrows(Exception.class, - () -> invokePrivate("validateConnection", new Class[]{})); - assertTrue(ex.getCause().getMessage().contains("Failed to connect")); - } - } - - // ========== Tests: concurrent map state ========== - - @Test - void testConcurrentMapState_multipleModels() throws Exception { - getModelStates().put("m1", RemoteModelStateEnum.ACTIVE); - getModelStates().put("m2", RemoteModelStateEnum.WARMING); - getModelStates().put("m3", RemoteModelStateEnum.COLD); - getModelStates().put("m4", RemoteModelStateEnum.FAILED); - - getModelClusterIps().put("m1", "10.0.0.1"); - getModelClusterIps().put("m2", "10.0.0.2"); - - getModelNames().put("m1", "Model-A"); - getModelNames().put("m2", "Model-B"); - getModelNames().put("m3", "Model-C"); - - assertEquals(RemoteModelStateEnum.ACTIVE, proxy.getModelState("m1")); - assertEquals(RemoteModelStateEnum.WARMING, proxy.getModelState("m2")); - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("m3")); - assertEquals(RemoteModelStateEnum.FAILED, proxy.getModelState("m4")); - - assertEquals("10.0.0.1", proxy.getModelClusterIp("m1")); - assertEquals("10.0.0.2", proxy.getModelClusterIp("m2")); - - assertEquals("Model-A", proxy.getModelName("m1")); - assertEquals("Model-B", proxy.getModelName("m2")); - assertEquals("Model-C", proxy.getModelName("m3")); - assertNull(proxy.getModelName("m4")); - } - - @Test - void testModelStateOverwrite() throws Exception { - getModelStates().put("m1", RemoteModelStateEnum.WARMING); - assertEquals(RemoteModelStateEnum.WARMING, proxy.getModelState("m1")); - - getModelStates().put("m1", RemoteModelStateEnum.ACTIVE); - assertEquals(RemoteModelStateEnum.ACTIVE, proxy.getModelState("m1")); - - getModelStates().put("m1", RemoteModelStateEnum.COLD); - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("m1")); - } - - // ========== Tests: loadInitialState (private) ========== - - @Test - void testLoadInitialState_loadsActiveAndWarming() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - - // First call: getZNodeData for MODEL_SCALER_PATH - JSONObject scalerResponse = new JSONObject(); - scalerResponse.put("data", "10.0.0.100"); - CloseableHttpResponse scalerResp = createMockResponse(200, scalerResponse.toString()); - - // Second call: getZNodeChildren for ACTIVE_PATH - JSONObject activeChildren = new JSONObject(); - activeChildren.put("children", new org.json.JSONArray(List.of("a1"))); - CloseableHttpResponse activeChildrenResp = createMockResponse(200, activeChildren.toString()); - - // Third call: getZNodeDataAsJson for a1 - JSONObject a1Data = new JSONObject(); - a1Data.put("ip", "10.0.0.10"); - a1Data.put("model_name", "ActiveModel"); - JSONObject a1Response = new JSONObject(); - a1Response.put("data", a1Data); - CloseableHttpResponse a1DataResp = createMockResponse(200, a1Response.toString()); - - // Fourth call: getZNodeChildren for WARMING_PATH - JSONObject warmingChildren = new JSONObject(); - warmingChildren.put("children", new org.json.JSONArray(List.of("w1"))); - CloseableHttpResponse warmingChildrenResp = createMockResponse(200, warmingChildren.toString()); - - when(mockClient.execute(any(HttpGet.class))) - .thenReturn(scalerResp) - .thenReturn(activeChildrenResp) - .thenReturn(a1DataResp) - .thenReturn(warmingChildrenResp); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - invokePrivate("loadInitialState", new Class[]{}); - - assertEquals(RemoteModelStateEnum.ACTIVE, getModelStates().get("a1")); - assertEquals("10.0.0.10", getModelClusterIps().get("a1")); - assertEquals("ActiveModel", getModelNames().get("a1")); - assertEquals(RemoteModelStateEnum.WARMING, getModelStates().get("w1")); - } - } - - @Test - void testLoadInitialState_httpError_doesNotThrow() throws Exception { - CloseableHttpClient mockClient = mock(CloseableHttpClient.class); - when(mockClient.execute(any(HttpGet.class))).thenThrow(new IOException("Network error")); - - try (MockedStatic mockedHttp = setupHttpMock(mockClient)) { - assertDoesNotThrow(() -> invokePrivate("loadInitialState", new Class[]{})); - } - } -} \ No newline at end of file diff --git a/test/prerna/cluster/util/RemoteClientServerZKUnitTests.java b/test/prerna/cluster/util/RemoteClientServerZKUnitTests.java deleted file mode 100644 index aed6e16a71d..00000000000 --- a/test/prerna/cluster/util/RemoteClientServerZKUnitTests.java +++ /dev/null @@ -1,839 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.*; -import static org.mockito.Mockito.*; - -import java.lang.reflect.Constructor; -import java.lang.reflect.Field; -import java.lang.reflect.Method; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.ScheduledExecutorService; - -import org.apache.curator.framework.CuratorFramework; -import org.apache.curator.framework.api.ExistsBuilder; -import org.apache.curator.framework.api.GetChildrenBuilder; -import org.apache.curator.framework.api.GetDataBuilder; -import org.apache.curator.framework.recipes.cache.CuratorCache; -import org.apache.zookeeper.data.Stat; -import org.json.JSONObject; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -import prerna.engine.api.RemoteModelStateEnum; - -class RemoteClientServerZKUnitTests { - - private RemoteClientServerZK proxy; - private CuratorFramework mockClient; - - @BeforeEach - void setUp() throws Exception { - Constructor ctor = - RemoteClientServerZK.class.getDeclaredConstructor(); - ctor.setAccessible(true); - proxy = ctor.newInstance(); - - mockClient = mock(CuratorFramework.class); - setField("client", mockClient); - } - - @AfterEach - void tearDown() throws Exception { - Field instanceField = RemoteClientServerZK.class.getDeclaredField("instance"); - instanceField.setAccessible(true); - instanceField.set(null, null); - } - - // ========== Reflection helpers ========== - - private void setField(String fieldName, Object value) throws Exception { - Field f = RemoteClientServerZK.class.getDeclaredField(fieldName); - f.setAccessible(true); - f.set(proxy, value); - } - - @SuppressWarnings("unchecked") - private ConcurrentMap getModelStates() throws Exception { - Field f = RemoteClientServerZK.class.getDeclaredField("modelStates"); - f.setAccessible(true); - return (ConcurrentMap) f.get(proxy); - } - - @SuppressWarnings("unchecked") - private ConcurrentMap getModelClusterIps() throws Exception { - Field f = RemoteClientServerZK.class.getDeclaredField("modelClusterIps"); - f.setAccessible(true); - return (ConcurrentMap) f.get(proxy); - } - - @SuppressWarnings("unchecked") - private ConcurrentMap getModelNames() throws Exception { - Field f = RemoteClientServerZK.class.getDeclaredField("modelNames"); - f.setAccessible(true); - return (ConcurrentMap) f.get(proxy); - } - - private ScheduledExecutorService getScheduler() throws Exception { - Field f = RemoteClientServerZK.class.getDeclaredField("scheduler"); - f.setAccessible(true); - return (ScheduledExecutorService) f.get(proxy); - } - - private Object invokePrivate(String methodName, Class[] paramTypes, Object... args) throws Exception { - Method m = RemoteClientServerZK.class.getDeclaredMethod(methodName, paramTypes); - m.setAccessible(true); - return m.invoke(proxy, args); - } - - // ========== Curator mock helpers ========== - - private void mockCheckExists(String path, boolean exists) throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath(path)).thenReturn(exists ? new Stat() : null); - } - - private void mockCheckExistsAny(boolean exists) throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath(anyString())).thenReturn(exists ? new Stat() : null); - } - - private void mockGetData(String path, byte[] data) throws Exception { - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath(path)).thenReturn(data); - } - - private void mockGetChildren(String path, List children) throws Exception { - GetChildrenBuilder getChildrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(getChildrenBuilder); - when(getChildrenBuilder.forPath(path)).thenReturn(children); - } - - private byte[] modelJsonBytes(String ip, String modelName) { - JSONObject json = new JSONObject(); - json.put("ip", ip); - json.put("model_name", modelName); - return json.toString().getBytes(StandardCharsets.UTF_8); - } - - // ========== Tests: getModelState ========== - - @Test - void testGetModelState_activeModel() throws Exception { - getModelStates().put("model-1", RemoteModelStateEnum.ACTIVE); - assertEquals(RemoteModelStateEnum.ACTIVE, proxy.getModelState("model-1")); - } - - @Test - void testGetModelState_warmingModel() throws Exception { - getModelStates().put("model-2", RemoteModelStateEnum.WARMING); - assertEquals(RemoteModelStateEnum.WARMING, proxy.getModelState("model-2")); - } - - @Test - void testGetModelState_coldModel() throws Exception { - getModelStates().put("model-3", RemoteModelStateEnum.COLD); - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("model-3")); - } - - @Test - void testGetModelState_failedModel() throws Exception { - getModelStates().put("model-4", RemoteModelStateEnum.FAILED); - assertEquals(RemoteModelStateEnum.FAILED, proxy.getModelState("model-4")); - } - - @Test - void testGetModelState_unknownModel_defaultsToCold() { - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("nonexistent")); - } - - @Test - void testGetModelState_emptyModelId() { - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("")); - } - - @Test - void testGetModelState_nullModelId_throwsNPE() { - assertThrows(NullPointerException.class, () -> proxy.getModelState(null)); - } - - // ========== Tests: getModelName ========== - - @Test - void testGetModelName_known() throws Exception { - getModelNames().put("model-1", "GPT-4"); - assertEquals("GPT-4", proxy.getModelName("model-1")); - } - - @Test - void testGetModelName_unknown() { - assertNull(proxy.getModelName("nonexistent")); - } - - @Test - void testGetModelName_emptyString() throws Exception { - getModelNames().put("model-1", ""); - assertEquals("", proxy.getModelName("model-1")); - } - - @Test - void testGetModelName_multipleModels() throws Exception { - getModelNames().put("m1", "Model-A"); - getModelNames().put("m2", "Model-B"); - getModelNames().put("m3", "Model-C"); - assertEquals("Model-A", proxy.getModelName("m1")); - assertEquals("Model-B", proxy.getModelName("m2")); - assertEquals("Model-C", proxy.getModelName("m3")); - } - - // ========== Tests: getModelScalerIp ========== - - @Test - void testGetModelScalerIp_set() throws Exception { - setField("modelScalerIp", "10.0.0.100"); - assertEquals("10.0.0.100", proxy.getModelScalerIp()); - } - - @Test - void testGetModelScalerIp_null() { - assertNull(proxy.getModelScalerIp()); - } - - // ========== Tests: getModelClusterIp ========== - - @Test - void testGetModelClusterIp_success() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/models/active/m1")).thenReturn(new Stat()); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/models/active/m1")).thenReturn(modelJsonBytes("10.0.0.5", "TestModel")); - - String ip = proxy.getModelClusterIp("m1"); - assertEquals("10.0.0.5", ip); - assertEquals("10.0.0.5", getModelClusterIps().get("m1")); - assertEquals("TestModel", getModelNames().get("m1")); - } - - @Test - void testGetModelClusterIp_pathDoesNotExist() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/models/active/m1")).thenReturn(null); - - assertNull(proxy.getModelClusterIp("m1")); - } - - @Test - void testGetModelClusterIp_noData() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/models/active/m1")).thenReturn(new Stat()); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/models/active/m1")).thenReturn(null); - - assertNull(proxy.getModelClusterIp("m1")); - } - - @Test - void testGetModelClusterIp_emptyData() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/models/active/m1")).thenReturn(new Stat()); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/models/active/m1")).thenReturn(new byte[0]); - - assertNull(proxy.getModelClusterIp("m1")); - } - - @Test - void testGetModelClusterIp_exception() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath(anyString())).thenThrow(new RuntimeException("ZK error")); - - assertNull(proxy.getModelClusterIp("m1")); - } - - @Test - void testGetModelClusterIp_updatesCache() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/models/active/m1")).thenReturn(new Stat()); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/models/active/m1")).thenReturn(modelJsonBytes("10.0.0.99", "CachedModel")); - - proxy.getModelClusterIp("m1"); - assertEquals("10.0.0.99", getModelClusterIps().get("m1")); - assertEquals("CachedModel", getModelNames().get("m1")); - } - - // ========== Tests: isModelWarming ========== - - @Test - void testIsModelWarming_true() throws Exception { - mockCheckExists("/models/warming/m1", true); - assertTrue(proxy.isModelWarming("m1")); - } - - @Test - void testIsModelWarming_false() throws Exception { - mockCheckExists("/models/warming/m1", false); - assertFalse(proxy.isModelWarming("m1")); - } - - @Test - void testIsModelWarming_exception() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath(anyString())).thenThrow(new RuntimeException("ZK error")); - assertFalse(proxy.isModelWarming("m1")); - } - - // ========== Tests: isModelActive ========== - - @Test - void testIsModelActive_true() throws Exception { - mockCheckExists("/models/active/m1", true); - assertTrue(proxy.isModelActive("m1")); - } - - @Test - void testIsModelActive_false() throws Exception { - mockCheckExists("/models/active/m1", false); - assertFalse(proxy.isModelActive("m1")); - } - - @Test - void testIsModelActive_exception() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath(anyString())).thenThrow(new RuntimeException("ZK error")); - assertFalse(proxy.isModelActive("m1")); - } - - // ========== Tests: waitForModelActive ========== - - @Test - void testWaitForModelActive_alreadyActive() throws Exception { - mockCheckExistsAny(true); - assertTrue(proxy.waitForModelActive("m1", 5000)); - } - - @Test - void testWaitForModelActive_timeout() throws Exception { - mockCheckExistsAny(false); - assertFalse(proxy.waitForModelActive("m1", 100)); - } - - // ========== Tests: waitForState ========== - - @Test - void testWaitForState_alreadyInDesiredState() throws Exception { - getModelStates().put("m1", RemoteModelStateEnum.ACTIVE); - assertTrue(proxy.waitForState("m1", RemoteModelStateEnum.ACTIVE, 5000)); - } - - @Test - void testWaitForState_defaultsCold() { - assertTrue(proxy.waitForState("nonexistent", RemoteModelStateEnum.COLD, 5000)); - } - - @Test - void testWaitForState_warmingModel_timeout() throws Exception { - getModelStates().put("m1", RemoteModelStateEnum.WARMING); - assertFalse(proxy.waitForState("m1", RemoteModelStateEnum.ACTIVE, 100)); - } - - @Test - void testWaitForState_failedModel_earlyExit() throws Exception { - getModelStates().put("m1", RemoteModelStateEnum.FAILED); - assertFalse(proxy.waitForState("m1", RemoteModelStateEnum.ACTIVE, 5000)); - } - - @Test - void testWaitForState_waitingForWarming() throws Exception { - getModelStates().put("m1", RemoteModelStateEnum.WARMING); - assertTrue(proxy.waitForState("m1", RemoteModelStateEnum.WARMING, 5000)); - } - - // ========== Tests: getActiveModels ========== - - @Test - void testGetActiveModels_withCachedNames() throws Exception { - mockGetChildren("/models/active", List.of("m1", "m2")); - - getModelNames().put("m1", "Model-A"); - getModelNames().put("m2", "Model-B"); - getModelStates().put("m1", RemoteModelStateEnum.ACTIVE); - getModelStates().put("m2", RemoteModelStateEnum.ACTIVE); - - List result = proxy.getActiveModels(); - assertEquals(2, result.size()); - assertEquals("m1", result.get(0).getId()); - assertEquals("Model-A", result.get(0).getName()); - assertEquals(RemoteModelStateEnum.ACTIVE, result.get(0).getState()); - assertEquals("m2", result.get(1).getId()); - assertEquals("Model-B", result.get(1).getName()); - } - - @Test - void testGetActiveModels_empty() throws Exception { - mockGetChildren("/models/active", new ArrayList<>()); - - List result = proxy.getActiveModels(); - assertTrue(result.isEmpty()); - } - - @Test - void testGetActiveModels_unknownNameFallsBackToZK() throws Exception { - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(List.of("m1")); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/models/active/m1")).thenReturn(modelJsonBytes("10.0.0.1", "FetchedModel")); - - List result = proxy.getActiveModels(); - assertEquals(1, result.size()); - assertEquals("m1", result.get(0).getId()); - assertEquals("FetchedModel", result.get(0).getName()); - assertEquals("FetchedModel", getModelNames().get("m1")); - } - - @Test - void testGetActiveModels_unknownNameAndNullData() throws Exception { - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(List.of("m1")); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/models/active/m1")).thenReturn(null); - - List result = proxy.getActiveModels(); - // Model with null data is not added - assertTrue(result.isEmpty()); - } - - @Test - void testGetActiveModels_unknownNameAndEmptyData() throws Exception { - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(List.of("m1")); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/models/active/m1")).thenReturn(new byte[0]); - - List result = proxy.getActiveModels(); - assertTrue(result.isEmpty()); - } - - @Test - void testGetActiveModels_exception() throws Exception { - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath(anyString())).thenThrow(new RuntimeException("ZK error")); - - List result = proxy.getActiveModels(); - assertTrue(result.isEmpty()); - } - - @Test - void testGetActiveModels_stateDefaultsToCold() throws Exception { - mockGetChildren("/models/active", List.of("m1")); - getModelNames().put("m1", "Model-A"); - // no state in modelStates - - List result = proxy.getActiveModels(); - assertEquals(1, result.size()); - assertEquals(RemoteModelStateEnum.COLD, result.get(0).getState()); - } - - // ========== Tests: getWarmingModels ========== - - @Test - void testGetWarmingModels_withCachedNames() throws Exception { - mockGetChildren("/models/warming", List.of("w1", "w2")); - - getModelNames().put("w1", "Warming-A"); - getModelStates().put("w1", RemoteModelStateEnum.WARMING); - getModelStates().put("w2", RemoteModelStateEnum.WARMING); - - List result = proxy.getWarmingModels(); - assertEquals(2, result.size()); - assertEquals("w1", result.get(0).getId()); - assertEquals("Warming-A", result.get(0).getName()); - assertEquals(RemoteModelStateEnum.WARMING, result.get(0).getState()); - assertEquals("w2", result.get(1).getId()); - assertEquals("Warming...", result.get(1).getName()); - } - - @Test - void testGetWarmingModels_empty() throws Exception { - mockGetChildren("/models/warming", new ArrayList<>()); - - List result = proxy.getWarmingModels(); - assertTrue(result.isEmpty()); - } - - @Test - void testGetWarmingModels_unknownNameDefaultsToWarmingText() throws Exception { - mockGetChildren("/models/warming", List.of("w1")); - - List result = proxy.getWarmingModels(); - assertEquals(1, result.size()); - assertEquals("Warming...", result.get(0).getName()); - } - - @Test - void testGetWarmingModels_stateDefaultsToWarming() throws Exception { - mockGetChildren("/models/warming", List.of("w1")); - // No state in modelStates - - List result = proxy.getWarmingModels(); - assertEquals(1, result.size()); - assertEquals(RemoteModelStateEnum.WARMING, result.get(0).getState()); - } - - @Test - void testGetWarmingModels_exception() throws Exception { - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath(anyString())).thenThrow(new RuntimeException("ZK error")); - - List result = proxy.getWarmingModels(); - assertTrue(result.isEmpty()); - } - - // ========== Tests: close ========== - - @Test - void testClose_closesAllResources() throws Exception { - CuratorCache warmingCache = mock(CuratorCache.class); - CuratorCache activeCache = mock(CuratorCache.class); - setField("warmingCache", warmingCache); - setField("activeCache", activeCache); - - proxy.close(); - - verify(warmingCache).close(); - verify(activeCache).close(); - verify(mockClient).close(); - } - - @Test - void testClose_nullCaches() throws Exception { - setField("warmingCache", null); - setField("activeCache", null); - - assertDoesNotThrow(() -> proxy.close()); - verify(mockClient).close(); - } - - @Test - void testClose_nullClient() throws Exception { - setField("warmingCache", null); - setField("activeCache", null); - setField("client", null); - - assertDoesNotThrow(() -> proxy.close()); - } - - @Test - void testClose_exceptionInCacheClose() throws Exception { - CuratorCache warmingCache = mock(CuratorCache.class); - doThrow(new RuntimeException("close error")).when(warmingCache).close(); - setField("warmingCache", warmingCache); - setField("activeCache", null); - - // Should not propagate the exception - assertDoesNotThrow(() -> proxy.close()); - } - - // ========== Tests: getModelIdFromPath (private) ========== - - @Test - void testGetModelIdFromPath_standardPath() throws Exception { - String result = (String) invokePrivate("getModelIdFromPath", - new Class[]{String.class}, "/models/active/model-123"); - assertEquals("model-123", result); - } - - @Test - void testGetModelIdFromPath_warmingPath() throws Exception { - String result = (String) invokePrivate("getModelIdFromPath", - new Class[]{String.class}, "/models/warming/abc-def"); - assertEquals("abc-def", result); - } - - @Test - void testGetModelIdFromPath_singleSegment() throws Exception { - String result = (String) invokePrivate("getModelIdFromPath", - new Class[]{String.class}, "model-123"); - assertEquals("model-123", result); - } - - @Test - void testGetModelIdFromPath_trailingSlash() throws Exception { - String result = (String) invokePrivate("getModelIdFromPath", - new Class[]{String.class}, "/models/active/model-1/"); - // lastIndexOf('/') + 1 gives empty string when trailing slash - assertEquals("", result); - } - - @Test - void testGetModelIdFromPath_deepPath() throws Exception { - String result = (String) invokePrivate("getModelIdFromPath", - new Class[]{String.class}, "/a/b/c/d/model-99"); - assertEquals("model-99", result); - } - - // ========== Tests: getModelPath (private) ========== - - @Test - void testGetModelPath_returnsModelId() throws Exception { - String result = (String) invokePrivate("getModelPath", - new Class[]{String.class}, "model-123"); - assertEquals("model-123", result); - } - - @Test - void testGetModelPath_emptyString() throws Exception { - String result = (String) invokePrivate("getModelPath", - new Class[]{String.class}, ""); - assertEquals("", result); - } - - // ========== Tests: concurrent map state ========== - - @Test - void testConcurrentMapState_multipleModels() throws Exception { - getModelStates().put("m1", RemoteModelStateEnum.ACTIVE); - getModelStates().put("m2", RemoteModelStateEnum.WARMING); - getModelStates().put("m3", RemoteModelStateEnum.COLD); - getModelStates().put("m4", RemoteModelStateEnum.FAILED); - - getModelClusterIps().put("m1", "10.0.0.1"); - getModelClusterIps().put("m2", "10.0.0.2"); - - getModelNames().put("m1", "Model-A"); - getModelNames().put("m2", "Model-B"); - getModelNames().put("m3", "Model-C"); - - assertEquals(RemoteModelStateEnum.ACTIVE, proxy.getModelState("m1")); - assertEquals(RemoteModelStateEnum.WARMING, proxy.getModelState("m2")); - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("m3")); - assertEquals(RemoteModelStateEnum.FAILED, proxy.getModelState("m4")); - - assertEquals("Model-A", proxy.getModelName("m1")); - assertEquals("Model-B", proxy.getModelName("m2")); - assertEquals("Model-C", proxy.getModelName("m3")); - assertNull(proxy.getModelName("m4")); - } - - @Test - void testModelStateOverwrite() throws Exception { - getModelStates().put("m1", RemoteModelStateEnum.WARMING); - assertEquals(RemoteModelStateEnum.WARMING, proxy.getModelState("m1")); - - getModelStates().put("m1", RemoteModelStateEnum.ACTIVE); - assertEquals(RemoteModelStateEnum.ACTIVE, proxy.getModelState("m1")); - - getModelStates().put("m1", RemoteModelStateEnum.COLD); - assertEquals(RemoteModelStateEnum.COLD, proxy.getModelState("m1")); - } - - // ========== Tests: refreshModelStates (private) ========== - - @Test - void testRefreshModelStates_updatesStates() throws Exception { - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(List.of("m1")); - when(childrenBuilder.forPath("/models/warming")).thenReturn(List.of("m2")); - - // m2 needs an existing non-ACTIVE state for the warming check - getModelStates().put("m2", RemoteModelStateEnum.COLD); - - invokePrivate("refreshModelStates", new Class[]{}); - assertEquals(RemoteModelStateEnum.ACTIVE, getModelStates().get("m1")); - assertEquals(RemoteModelStateEnum.WARMING, getModelStates().get("m2")); - } - - @Test - void testRefreshModelStates_staleEntriesBecomeCold() throws Exception { - getModelStates().put("stale-model", RemoteModelStateEnum.ACTIVE); - - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(new ArrayList<>()); - when(childrenBuilder.forPath("/models/warming")).thenReturn(new ArrayList<>()); - - invokePrivate("refreshModelStates", new Class[]{}); - assertEquals(RemoteModelStateEnum.COLD, getModelStates().get("stale-model")); - } - - @Test - void testRefreshModelStates_activeOverridesWarming() throws Exception { - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(List.of("m1")); - when(childrenBuilder.forPath("/models/warming")).thenReturn(List.of("m1")); - - invokePrivate("refreshModelStates", new Class[]{}); - assertEquals(RemoteModelStateEnum.ACTIVE, getModelStates().get("m1")); - } - - // ========== Tests: loadInitialState (private) ========== - - @Test - void testLoadInitialState_loadsActiveAndWarming() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/services/kube-model-deployer")).thenReturn(new Stat()); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/services/kube-model-deployer")).thenReturn("10.0.0.100".getBytes(StandardCharsets.UTF_8)); - when(getDataBuilder.forPath("/models/active/a1")).thenReturn(modelJsonBytes("10.0.0.10", "ActiveModel")); - - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(List.of("a1")); - when(childrenBuilder.forPath("/models/warming")).thenReturn(List.of("w1")); - - invokePrivate("loadInitialState", new Class[]{}); - - assertEquals("10.0.0.100", proxy.getModelScalerIp()); - assertEquals(RemoteModelStateEnum.ACTIVE, getModelStates().get("a1")); - assertEquals("10.0.0.10", getModelClusterIps().get("a1")); - assertEquals("ActiveModel", getModelNames().get("a1")); - assertEquals(RemoteModelStateEnum.WARMING, getModelStates().get("w1")); - } - - @Test - void testLoadInitialState_noScalerPath() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/services/kube-model-deployer")).thenReturn(null); - - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(new ArrayList<>()); - when(childrenBuilder.forPath("/models/warming")).thenReturn(new ArrayList<>()); - - invokePrivate("loadInitialState", new Class[]{}); - assertNull(proxy.getModelScalerIp()); - } - - @Test - void testLoadInitialState_scalerPathEmptyData() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/services/kube-model-deployer")).thenReturn(new Stat()); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/services/kube-model-deployer")).thenReturn(new byte[0]); - - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(new ArrayList<>()); - when(childrenBuilder.forPath("/models/warming")).thenReturn(new ArrayList<>()); - - invokePrivate("loadInitialState", new Class[]{}); - assertNull(proxy.getModelScalerIp()); - } - - @Test - void testLoadInitialState_exception_doesNotThrow() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath(anyString())).thenThrow(new RuntimeException("ZK error")); - - assertDoesNotThrow(() -> invokePrivate("loadInitialState", new Class[]{})); - } - - @Test - void testLoadInitialState_activeModelWithNullData() throws Exception { - ExistsBuilder existsBuilder = mock(ExistsBuilder.class); - when(mockClient.checkExists()).thenReturn(existsBuilder); - when(existsBuilder.forPath("/services/kube-model-deployer")).thenReturn(null); - - GetDataBuilder getDataBuilder = mock(GetDataBuilder.class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath("/models/active/a1")).thenReturn(null); - - GetChildrenBuilder childrenBuilder = mock(GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(childrenBuilder); - when(childrenBuilder.forPath("/models/active")).thenReturn(List.of("a1")); - when(childrenBuilder.forPath("/models/warming")).thenReturn(new ArrayList<>()); - - invokePrivate("loadInitialState", new Class[]{}); - // a1 should not be in modelStates since data was null - assertNull(getModelStates().get("a1")); - } - - // ========== Tests: devPortForwarding field ========== - - @Test - void testDevPortForwarding_defaultFalse() throws Exception { - Field f = RemoteClientServerZK.class.getDeclaredField("devPortForwarding"); - f.setAccessible(true); - assertEquals(false, f.get(proxy)); - } - - // ========== Tests: zkServer field ========== - - @Test - void testZkServer_default() throws Exception { - Field f = RemoteClientServerZK.class.getDeclaredField("zkServer"); - f.setAccessible(true); - assertEquals("localhost:2181", f.get(proxy)); - } -} \ No newline at end of file diff --git a/test/prerna/cluster/util/RemoteModelInfoUnitTests.java b/test/prerna/cluster/util/RemoteModelInfoUnitTests.java deleted file mode 100644 index 8fa65aca2d5..00000000000 --- a/test/prerna/cluster/util/RemoteModelInfoUnitTests.java +++ /dev/null @@ -1,145 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; - -import org.junit.jupiter.api.Test; - -import prerna.engine.api.RemoteModelStateEnum; - -class RemoteModelInfoUnitTests { - - @Test - void testConstructorAndGetters() { - RemoteModelInfo info = new RemoteModelInfo("id1", "modelName", RemoteModelStateEnum.ACTIVE); - assertEquals("id1", info.getId()); - assertEquals("modelName", info.getName()); - assertEquals(RemoteModelStateEnum.ACTIVE, info.getState()); - } - - @Test - void testStateCold() { - RemoteModelInfo info = new RemoteModelInfo("c1", "cold-model", RemoteModelStateEnum.COLD); - assertEquals(RemoteModelStateEnum.COLD, info.getState()); - assertEquals("c1", info.getId()); - assertEquals("cold-model", info.getName()); - } - - @Test - void testStateWarming() { - RemoteModelInfo info = new RemoteModelInfo("w1", "warming-model", RemoteModelStateEnum.WARMING); - assertEquals(RemoteModelStateEnum.WARMING, info.getState()); - assertEquals("w1", info.getId()); - assertEquals("warming-model", info.getName()); - } - - @Test - void testStateActive() { - RemoteModelInfo info = new RemoteModelInfo("a1", "active-model", RemoteModelStateEnum.ACTIVE); - assertEquals(RemoteModelStateEnum.ACTIVE, info.getState()); - assertEquals("a1", info.getId()); - assertEquals("active-model", info.getName()); - } - - @Test - void testStateFailed() { - RemoteModelInfo info = new RemoteModelInfo("f1", "failed-model", RemoteModelStateEnum.FAILED); - assertEquals(RemoteModelStateEnum.FAILED, info.getState()); - assertEquals("f1", info.getId()); - assertEquals("failed-model", info.getName()); - } - - @Test - void testStateUnknown() { - RemoteModelInfo info = new RemoteModelInfo("u1", "unknown-model", RemoteModelStateEnum.UNKNOWN); - assertEquals(RemoteModelStateEnum.UNKNOWN, info.getState()); - assertEquals("u1", info.getId()); - assertEquals("unknown-model", info.getName()); - } - - @Test - void testNullId() { - RemoteModelInfo info = new RemoteModelInfo(null, "name", RemoteModelStateEnum.ACTIVE); - assertNull(info.getId()); - assertEquals("name", info.getName()); - assertEquals(RemoteModelStateEnum.ACTIVE, info.getState()); - } - - @Test - void testNullName() { - RemoteModelInfo info = new RemoteModelInfo("id", null, RemoteModelStateEnum.ACTIVE); - assertEquals("id", info.getId()); - assertNull(info.getName()); - assertEquals(RemoteModelStateEnum.ACTIVE, info.getState()); - } - - @Test - void testNullState() { - RemoteModelInfo info = new RemoteModelInfo("id", "name", null); - assertEquals("id", info.getId()); - assertEquals("name", info.getName()); - assertNull(info.getState()); - } - - @Test - void testAllNulls() { - RemoteModelInfo info = new RemoteModelInfo(null, null, null); - assertNull(info.getId()); - assertNull(info.getName()); - assertNull(info.getState()); - } - - @Test - void testEmptyStringId() { - RemoteModelInfo info = new RemoteModelInfo("", "name", RemoteModelStateEnum.COLD); - assertEquals("", info.getId()); - assertEquals("name", info.getName()); - } - - @Test - void testEmptyStringName() { - RemoteModelInfo info = new RemoteModelInfo("id", "", RemoteModelStateEnum.COLD); - assertEquals("id", info.getId()); - assertEquals("", info.getName()); - } - - @Test - void testEmptyStringsIdAndName() { - RemoteModelInfo info = new RemoteModelInfo("", "", RemoteModelStateEnum.WARMING); - assertEquals("", info.getId()); - assertEquals("", info.getName()); - assertEquals(RemoteModelStateEnum.WARMING, info.getState()); - } - - @Test - void testAllEnumValuesExist() { - RemoteModelStateEnum[] values = RemoteModelStateEnum.values(); - assertEquals(5, values.length, "Expected exactly 5 enum values"); - } -} diff --git a/test/prerna/cluster/util/ZKClientFactoryUnitTests.java b/test/prerna/cluster/util/ZKClientFactoryUnitTests.java deleted file mode 100644 index 49e72b4811a..00000000000 --- a/test/prerna/cluster/util/ZKClientFactoryUnitTests.java +++ /dev/null @@ -1,77 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.cluster.util; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -import org.junit.jupiter.api.Test; -import org.mockito.MockedStatic; - -class ZKClientFactoryUnitTests { - - @Test - void testGetZKClient_devPortForwardingTrue_returnsZKInstance() { - RemoteClientServerZK expectedClient = mock(RemoteClientServerZK.class); - try (MockedStatic zkMockedStatic = mockStatic(RemoteClientServerZK.class)) { - zkMockedStatic.when(RemoteClientServerZK::getInstance).thenReturn(expectedClient); - IRemoteClientServer result = ZKClientFactory.getZKClient(Boolean.TRUE); - assertSame(expectedClient, result); - } - } - - @Test - void testGetZKClient_zkIngressNotSet_returnsZKInstance() { - RemoteClientServerZK expectedClient = mock(RemoteClientServerZK.class); - try (MockedStatic zkMockedStatic = mockStatic(RemoteClientServerZK.class)) { - zkMockedStatic.when(RemoteClientServerZK::getInstance).thenReturn(expectedClient); - IRemoteClientServer result = ZKClientFactory.getZKClient(Boolean.FALSE); - assertSame(expectedClient, result); - } - } - - @Test - void testGetZKClient_devPortForwardingFalse_noZkIngress_returnsZKInstance() { - RemoteClientServerZK expectedClient = mock(RemoteClientServerZK.class); - try (MockedStatic zkMockedStatic = mockStatic(RemoteClientServerZK.class)) { - zkMockedStatic.when(RemoteClientServerZK::getInstance).thenReturn(expectedClient); - IRemoteClientServer result = ZKClientFactory.getZKClient(false); - assertSame(expectedClient, result); - } - } - - @Test - void testGetZKClient_resultNotNull() { - RemoteClientServerZK expectedClient = mock(RemoteClientServerZK.class); - try (MockedStatic zkMockedStatic = mockStatic(RemoteClientServerZK.class)) { - zkMockedStatic.when(RemoteClientServerZK::getInstance).thenReturn(expectedClient); - IRemoteClientServer result = ZKClientFactory.getZKClient(true); - assertNotNull(result); - } - } -} diff --git a/test/prerna/engine/api/ModelTypeEnumUnitTests.java b/test/prerna/engine/api/ModelTypeEnumUnitTests.java index 83926207cea..e084cdaac32 100644 --- a/test/prerna/engine/api/ModelTypeEnumUnitTests.java +++ b/test/prerna/engine/api/ModelTypeEnumUnitTests.java @@ -36,9 +36,6 @@ import prerna.engine.api.ModelTypeEnum; import prerna.engine.impl.model.BedrockEngine; import prerna.engine.impl.model.EmbeddedModelEngine; -import prerna.engine.impl.model.KServeImageEmbedEngine; -import prerna.engine.impl.model.KServeVisionEngine; -import prerna.engine.impl.model.NEREngine; import prerna.engine.impl.model.OpenAiEngine; import prerna.engine.impl.model.TextEmbeddingsEngine; import prerna.engine.impl.model.TextGenerationEngine; @@ -61,27 +58,6 @@ void testEmbedded() { assertEquals(EmbeddedModelEngine.class.getName(), testEnum.getModelClass()); } - @Test - void testNER() { - ModelTypeEnum testEnum = ModelTypeEnum.NER; - assertEquals("NER", testEnum.getModelName()); - assertEquals(NEREngine.class.getName(), testEnum.getModelClass()); - } - - @Test - void testVision() { - ModelTypeEnum testEnum = ModelTypeEnum.KSERVE_VISION; - assertEquals("KSERVE_VISION", testEnum.getModelName()); - assertEquals(KServeVisionEngine.class.getName(), testEnum.getModelClass()); - } - - @Test - void testImageEmbed() { - ModelTypeEnum testEnum = ModelTypeEnum.KSERVE_IMAGE_EMBED; - assertEquals("KSERVE_IMAGE_EMBED", testEnum.getModelName()); - assertEquals(KServeImageEmbedEngine.class.getName(), testEnum.getModelClass()); - } - @Test void testOpenAI() { ModelTypeEnum testEnum = ModelTypeEnum.OPEN_AI; diff --git a/test/prerna/engine/api/RemoteModelStateEnumUnitTests.java b/test/prerna/engine/api/RemoteModelStateEnumUnitTests.java deleted file mode 100644 index b363c4b8986..00000000000 --- a/test/prerna/engine/api/RemoteModelStateEnumUnitTests.java +++ /dev/null @@ -1,65 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.api; - -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.Arrays; -import java.util.List; -import java.util.Vector; -import java.util.stream.Collectors; -import java.util.stream.Stream; - -import org.apache.commons.collections4.CollectionUtils; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -import prerna.engine.api.RemoteModelStateEnum; - -public class RemoteModelStateEnumUnitTests { - private RemoteModelStateEnum[] enumArr; - - @BeforeEach - void setUp() { - enumArr = RemoteModelStateEnum.values(); - } - - @Test - void testAllEnumsExist() { - List enumLst = Stream.of( - RemoteModelStateEnum.COLD, - RemoteModelStateEnum.WARMING, - RemoteModelStateEnum.ACTIVE, - RemoteModelStateEnum.FAILED, - RemoteModelStateEnum.UNKNOWN - ).collect(Collectors.toCollection(Vector::new)); - - assertTrue(CollectionUtils.containsAll(enumLst, Arrays.asList(enumArr))); - } - -} diff --git a/test/prerna/engine/impl/model/kserve/KServeAdapterUnitTests.java b/test/prerna/engine/impl/model/kserve/KServeAdapterUnitTests.java deleted file mode 100644 index 82954782688..00000000000 --- a/test/prerna/engine/impl/model/kserve/KServeAdapterUnitTests.java +++ /dev/null @@ -1,126 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model.kserve; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Set; - -import org.json.JSONObject; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - - -import prerna.engine.impl.model.kserve.KServeAdapter; - -public class KServeAdapterUnitTests { - private KServeAdapter classToTest; - - @BeforeEach - public void setUp() { - classToTest = new KServeAdapter(); - } - - @Test - public void executeRequest() throws Exception{ - JSONObject jsonObj = new JSONObject(); - jsonObj.put("model", ""); - jsonObj.put("string", "value"); - jsonObj.put("JSONArray", new ArrayList(){{add(1);}}); - - JSONObject retVal = classToTest.toKServeRequest(jsonObj); - - assertNotNull(retVal); - assertEquals(2, retVal.getJSONArray("inputs").length()); - - assertEquals(1, ((int[]) retVal.getJSONArray("inputs").getJSONObject(0).get("shape"))[0]); - assertEquals("value", retVal.getJSONArray("inputs").getJSONObject(0).getJSONArray("data").getString(0)); - assertEquals("BYTES", retVal.getJSONArray("inputs").getJSONObject(0).get("datatype")); - assertEquals("string", retVal.getJSONArray("inputs").getJSONObject(0).get("name")); - - assertEquals(1, ((int[]) retVal.getJSONArray("inputs").getJSONObject(1).get("shape"))[0]); - assertEquals(1, retVal.getJSONArray("inputs").getJSONObject(1).getJSONArray("data").getInt(0)); - assertEquals("BYTES", retVal.getJSONArray("inputs").getJSONObject(1).get("datatype")); - assertEquals("JSONArray", retVal.getJSONArray("inputs").getJSONObject(1).get("name")); - } - - @Test - public void executeRequestList() throws Exception{ - JSONObject mock = mock(JSONObject.class); - - Set set = new HashSet<>(); - set.add("test"); - when(mock.keySet()).thenReturn(set); - - List list = new ArrayList<>(); - list.add("value"); - when(mock.get("test")).thenReturn(list); - - JSONObject retVal = classToTest.toKServeRequest(mock); - - assertNotNull(retVal); - assertEquals(1, retVal.getJSONArray("inputs").length()); - assertEquals(1, ((int[]) retVal.getJSONArray("inputs").getJSONObject(0).get("shape"))[0]); - assertEquals("value", retVal.getJSONArray("inputs").getJSONObject(0).getJSONArray("data").getString(0)); - assertEquals("BYTES", retVal.getJSONArray("inputs").getJSONObject(0).get("datatype")); - assertEquals("test", retVal.getJSONArray("inputs").getJSONObject(0).get("name")); - } - - @Test - public void executeResponse() { - ArrayList list = new ArrayList<>(); - - list.add(new JSONObject("{name:nameVal1, data:[]}")); - list.add(new JSONObject("{name:nameVal2, data:[val2]}")); - list.add(new JSONObject("{name:nameVal3, data:[\"{val3}\"]}")); - list.add(new JSONObject("{name:nameVal4, data:[\"{val4:1}\"]}")); - list.add(new JSONObject("{name:nameVal5, data:[\"[val5]\"]}")); - list.add(new JSONObject("{name:nameVal6, data:['val6', 'val']}")); - - JSONObject kserveResponse = new JSONObject(); - kserveResponse.put("outputs", list); - - JSONObject retVal = classToTest.formatKServeResponse(kserveResponse); - - assertNotNull(retVal); - assertEquals(6, retVal.length()); - - assertEquals("[]", retVal.get("nameVal1").toString()); - assertEquals("val2", retVal.get("nameVal2").toString()); - assertEquals("{val3}", retVal.get("nameVal3").toString()); - assertEquals("{\"val4\":1}", retVal.get("nameVal4").toString()); - assertEquals("[\"val5\"]", retVal.get("nameVal5").toString()); - assertEquals("[\"val6\",\"val\"]", retVal.get("nameVal6").toString()); - } -} diff --git a/test/prerna/engine/impl/model/responses/AskImageModelEngineResponseUnitTests.java b/test/prerna/engine/impl/model/responses/AskImageModelEngineResponseUnitTests.java index 0987220e21f..4640f6458dd 100644 --- a/test/prerna/engine/impl/model/responses/AskImageModelEngineResponseUnitTests.java +++ b/test/prerna/engine/impl/model/responses/AskImageModelEngineResponseUnitTests.java @@ -41,99 +41,6 @@ public class AskImageModelEngineResponseUnitTests { private AskImageModelEngineResponse reactor; - @Test - void getKServeOutputElse() { - Map map = new HashMap(); - map.put("string", "string"); - - JSONObject jsonObj = new JSONObject(); - jsonObj.put("output", map); - jsonObj.put("prompt", "prompt"); - jsonObj.put("negative_prompt", "negative_prompt"); - jsonObj.put("height", 1); - jsonObj.put("width", 1); - jsonObj.put("num_inference_steps", 1); - jsonObj.put("guidance_scale", 1.0); - jsonObj.put("seed", 1); - - reactor = new AskImageModelEngineResponse(null, 0, 0); - AskImageModelEngineResponse obj = reactor.getKServeImageResponse(jsonObj); - - assertNotNull(obj); - - assertNotNull(obj.getResponse().get("images")); - assertEquals("{\"string\":\"string\"}", obj.getResponse().get("images").toString()); - assertEquals(0, (int) obj.getNumberOfTokensInPrompt()); - assertEquals(0,(int) obj.getNumberOfTokensInResponse()); - - assertEquals("prompt", obj.getResponse().get("prompt")); - assertEquals("negative_prompt", obj.getResponse().get("negative_prompt")); - assertEquals(1, obj.getResponse().get("height")); - assertEquals(1, obj.getResponse().get("width")); - assertEquals(1, obj.getResponse().get("num_inference_steps")); - assertEquals(1.0, obj.getResponse().get("guidance_scale")); - assertEquals(1, obj.getResponse().get("seed")); - } - - @Test - void getKServeOutputJsonArray() { - ArrayList list = new ArrayList(); - list.add("JsonArray"); - JSONObject jsonObj = new JSONObject(); - jsonObj.put("output", list); - - reactor = new AskImageModelEngineResponse(null, 0, 0); - AskImageModelEngineResponse obj = reactor.getKServeImageResponse(jsonObj); - - assertNotNull(obj); - assertNotNull(obj.getResponse().get("images")); - assertEquals("JsonArray", ((String[])obj.getResponse().get("images"))[0]); - assertEquals(0, (int) obj.getNumberOfTokensInPrompt()); - assertEquals(0, (int) obj.getNumberOfTokensInResponse()); - - } - - @Test - void getKserverOutputStringJsonArray() { - JSONObject jsonObj = new JSONObject(); - jsonObj.put("output", "[string]"); - - reactor = new AskImageModelEngineResponse(null, 0, 0); - AskImageModelEngineResponse obj = reactor.getKServeImageResponse(jsonObj); - - assertNotNull(obj); - assertNotNull(obj.getResponse().get("images")); - assertEquals("string", ((String[])obj.getResponse().get("images"))[0]); - assertEquals(0, (int) obj.getNumberOfTokensInPrompt()); - assertEquals(0, (int) obj.getNumberOfTokensInResponse()); - - } - - @Test - void getKserverOutputStringCatch() { - JSONObject jsonObj = new JSONObject(); - jsonObj.put("output", "[{json:obj}]"); - - reactor = new AskImageModelEngineResponse(null, 0, 0); - AskImageModelEngineResponse obj = reactor.getKServeImageResponse(jsonObj); - - assertNotNull(obj); - assertNotNull(obj.getResponse().get("images")); - assertEquals("[{json:obj}]", obj.getResponse().get("images").toString()); - } - - @Test - void getKserverNullParam() { - reactor = new AskImageModelEngineResponse(null, 0, 0); - AskImageModelEngineResponse obj = reactor.getKServeImageResponse(null); - - assertNotNull(obj); - assertEquals("error", obj.getResponse().get("status")); - assertEquals("Null response from model request", obj.getResponse().get("message")); - assertEquals(0, (int) obj.getNumberOfTokensInPrompt()); - assertEquals(0, (int) obj.getNumberOfTokensInResponse()); - } - @Test void getStringImagesJsonArray() { Map map = new HashMap(); diff --git a/test/prerna/engine/impl/model/responses/NerModelEngineResponseUnitTests.java b/test/prerna/engine/impl/model/responses/NerModelEngineResponseUnitTests.java deleted file mode 100644 index 12088340943..00000000000 --- a/test/prerna/engine/impl/model/responses/NerModelEngineResponseUnitTests.java +++ /dev/null @@ -1,107 +0,0 @@ -/******************************************************************************* - * Copyright 2015 Defense Health Agency (DHA) - * - * If your use of this software does not include any GPLv2 components: - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * ---------------------------------------------------------------------------- - * If your use of this software includes any GPLv2 components: - * This program is free software; you can redistribute it and/or - * modify it under the terms of the GNU General Public License - * as published by the Free Software Foundation; either version 2 - * of the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - *******************************************************************************/ -package prerna.engine.impl.model.responses; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; - -import java.util.HashMap; -import java.util.Map; - -import org.json.JSONArray; -import org.json.JSONObject; -import org.junit.jupiter.api.Test; - -public class NerModelEngineResponseUnitTests { - private NerModelEngineResponse reactor; - - @Test - void test() { - Map map = new HashMap(){{put("key", "value");}}; - - NerModelEngineResponse ans = new NerModelEngineResponse(map, 1, 2); - ans.setMessageId("messageId"); - ans.setRoomId("roomId"); - - assertNotNull(ans); - assertEquals("{key=value}", ans.getResponse().toString()); - assertEquals("messageId", ans.getMessageId()); - assertEquals("roomId", ans.getRoomId()); - assertEquals(1, (int) ans.getNumberOfTokensInPrompt()); - assertEquals(2, (int) ans.getNumberOfTokensInResponse()); - } - - @Test - void fromJson() { - JSONObject obj = new JSONObject(); - obj.put("mask_values", new JSONObject(){{put("key","value");}}); - obj.put("entities", new JSONArray(){{put("string");}}); - obj.put("raw_output", new JSONArray(){{ - put(new JSONObject(){{ - put("start", 0); - put("end", 1); - put("text", "text"); - put("label", "label"); - put("score", 5.0); - }}); - }}); - obj.put("output", "output"); - obj.put("input", "input"); - obj.put("status", "status"); - obj.put("message", "message"); - - NerModelEngineResponse ans = reactor.fromJson(obj); - - assertNotNull(ans); - assertEquals("{mask_values={key=value}, output=output, input=input, entities=[string], raw_output=[{score=5.0, start=0, end=1, text=text, label=label}], message=message, status=status}", ans.getResponse().toString()); - assertEquals(0, (int) ans.getNumberOfTokensInPrompt()); - assertEquals(0, (int) ans.getNumberOfTokensInResponse()); - } - - @Test - void fromJsonEmptyObj() { - JSONObject obj = new JSONObject(); - - NerModelEngineResponse ans = reactor.fromJson(obj); - - assertNotNull(ans); - assertEquals("{message=, status=success}", ans.getResponse().toString()); - assertEquals(0, (int) ans.getNumberOfTokensInPrompt()); - assertEquals(0, (int) ans.getNumberOfTokensInResponse()); - } - - @Test - void fromJsonNullObj() { - NerModelEngineResponse ans = reactor.fromJson(null); - - assertNotNull(ans); - assertEquals("{message=Null response from model request, status=error}", ans.getResponse().toString()); - assertEquals(0, (int) ans.getNumberOfTokensInPrompt()); - assertEquals(0, (int) ans.getNumberOfTokensInResponse()); - } -} diff --git a/test/prerna/reactor/automation/run/AutomationPixelExpressionRendererUnitTests.java b/test/prerna/reactor/automation/run/AutomationPixelExpressionRendererUnitTests.java index 2b90c7e6307..ba782ba41f2 100644 --- a/test/prerna/reactor/automation/run/AutomationPixelExpressionRendererUnitTests.java +++ b/test/prerna/reactor/automation/run/AutomationPixelExpressionRendererUnitTests.java @@ -227,11 +227,7 @@ void rendersCompleteFrameRowsInExactAndEmbeddedModelPrompts() throws Exception { } @Test - void rejectsModelErrorAndMissingVisionMedia() { - assertThrows(IllegalStateException.class, - () -> AutomationPixelExpressionRenderer.modelResult( - prerna.reactor.automation.AutomationConstants.NODE_MODEL_NER, - Map.of("response", Map.of("status", "error", "message", "failed")), "ner")); + void rejectsMissingVisionMedia() { assertThrows(IllegalArgumentException.class, () -> AutomationPixelExpressionRenderer.model( prerna.reactor.automation.AutomationConstants.NODE_MODEL_VISION,