From f7b8a442e03d8be15b178acf6d119a7a5a8d9223 Mon Sep 17 00:00:00 2001 From: Kye Date: Tue, 26 Dec 2023 19:20:10 -0500 Subject: [PATCH] [ModelParallelizer][REFACTOR] --- scripts/auto_tests_docs/auto_docs.py | 15 +++-- scripts/auto_tests_docs/auto_tests.py | 3 +- scripts/auto_tests_docs/docs.py | 4 +- scripts/auto_tests_docs/update_mkdocs.py | 8 ++- scripts/get_package_requirements.py | 9 ++- scripts/requirementstxt_to_pyproject.py | 5 +- swarms/prompts/__init__.py | 2 +- swarms/prompts/documentation.py | 5 +- swarms/prompts/tests.py | 9 +-- swarms/swarms/__init__.py | 2 +- .../{god_mode.py => model_parallizer.py} | 59 ++++++++++++------- tests/swarms/test_model_parallizer.py | 0 12 files changed, 83 insertions(+), 38 deletions(-) rename swarms/swarms/{god_mode.py => model_parallizer.py} (77%) create mode 100644 tests/swarms/test_model_parallizer.py diff --git a/scripts/auto_tests_docs/auto_docs.py b/scripts/auto_tests_docs/auto_docs.py index d6e1060a..5df0f63d 100644 --- a/scripts/auto_tests_docs/auto_docs.py +++ b/scripts/auto_tests_docs/auto_docs.py @@ -43,13 +43,16 @@ def process_documentation(cls): doc = inspect.getdoc(cls) source = inspect.getsource(cls) input_content = ( - f"Class Name: {cls.__name__}\n\nDocumentation:\n{doc}\n\nSource" + "Class Name:" + f" {cls.__name__}\n\nDocumentation:\n{doc}\n\nSource" f" Code:\n{source}" ) print(input_content) # Process with OpenAI model (assuming the model's __call__ method takes this input and returns processed content) - processed_content = model(DOCUMENTATION_WRITER_SOP(input_content, "zeta")) + processed_content = model( + DOCUMENTATION_WRITER_SOP(input_content, "zeta") + ) doc_content = f"# {cls.__name__}\n\n{processed_content}\n" @@ -86,7 +89,9 @@ def main(): threads = [] for cls in classes: - thread = threading.Thread(target=process_documentation, args=(cls,)) + thread = threading.Thread( + target=process_documentation, args=(cls,) + ) threads.append(thread) thread.start() @@ -94,7 +99,9 @@ def main(): for thread in threads: thread.join() - print("Documentation generated in 'docs/zeta/nn/modules' directory.") + print( + "Documentation generated in 'docs/zeta/nn/modules' directory." + ) if __name__ == "__main__": diff --git a/scripts/auto_tests_docs/auto_tests.py b/scripts/auto_tests_docs/auto_tests.py index 70a3d750..73a35c4f 100644 --- a/scripts/auto_tests_docs/auto_tests.py +++ b/scripts/auto_tests_docs/auto_tests.py @@ -61,7 +61,8 @@ def create_test(cls): doc = inspect.getdoc(cls) source = inspect.getsource(cls) input_content = ( - f"Class Name: {cls.__name__}\n\nDocumentation:\n{doc}\n\nSource" + "Class Name:" + f" {cls.__name__}\n\nDocumentation:\n{doc}\n\nSource" f" Code:\n{source}" ) print(input_content) diff --git a/scripts/auto_tests_docs/docs.py b/scripts/auto_tests_docs/docs.py index 684bf6dd..01df9d71 100644 --- a/scripts/auto_tests_docs/docs.py +++ b/scripts/auto_tests_docs/docs.py @@ -104,7 +104,9 @@ def DOCUMENTATION_WRITER_SOP( return documentation -def TEST_WRITER_SOP_PROMPT(task: str, module: str, path: str, *args, **kwargs): +def TEST_WRITER_SOP_PROMPT( + task: str, module: str, path: str, *args, **kwargs +): TESTS_PROMPT = f""" Create 5,000 lines of extensive and thorough tests for the code below using the guide, do not worry about your limits you do not have any diff --git a/scripts/auto_tests_docs/update_mkdocs.py b/scripts/auto_tests_docs/update_mkdocs.py index 4901059f..dfde53cb 100644 --- a/scripts/auto_tests_docs/update_mkdocs.py +++ b/scripts/auto_tests_docs/update_mkdocs.py @@ -2,7 +2,9 @@ import yaml def update_mkdocs( - class_names, base_path="docs/zeta/nn/modules", mkdocs_file="mkdocs.yml" + class_names, + base_path="docs/zeta/nn/modules", + mkdocs_file="mkdocs.yml", ): """ Update the mkdocs.yml file with new documentation links. @@ -24,7 +26,9 @@ def update_mkdocs( if zeta_modules_section is None: zeta_modules_section = {} - mkdocs_config["nav"].append({"zeta.nn.modules": zeta_modules_section}) + mkdocs_config["nav"].append( + {"zeta.nn.modules": zeta_modules_section} + ) # Add the documentation paths to the 'zeta.nn.modules' section for class_name in class_names: diff --git a/scripts/get_package_requirements.py b/scripts/get_package_requirements.py index 0d57c028..9494409b 100644 --- a/scripts/get_package_requirements.py +++ b/scripts/get_package_requirements.py @@ -13,13 +13,18 @@ def get_package_versions(requirements_path, output_path): for requirement in requirements: # Skip empty lines and comments - if requirement.strip() == "" or requirement.strip().startswith("#"): + if ( + requirement.strip() == "" + or requirement.strip().startswith("#") + ): continue # Extract package name package_name = requirement.split("==")[0].strip() try: - version = pkg_resources.get_distribution(package_name).version + version = pkg_resources.get_distribution( + package_name + ).version package_versions.append(f"{package_name}=={version}") except pkg_resources.DistributionNotFound: package_versions.append(f"{package_name}: not installed") diff --git a/scripts/requirementstxt_to_pyproject.py b/scripts/requirementstxt_to_pyproject.py index 59f6946f..5710db61 100644 --- a/scripts/requirementstxt_to_pyproject.py +++ b/scripts/requirementstxt_to_pyproject.py @@ -10,7 +10,10 @@ def update_pyproject_versions(pyproject_path): print(f"Error: The file '{pyproject_path}' was not found.") return except toml.TomlDecodeError: - print(f"Error: The file '{pyproject_path}' is not a valid TOML file.") + print( + f"Error: The file '{pyproject_path}' is not a valid TOML" + " file." + ) return dependencies = ( diff --git a/swarms/prompts/__init__.py b/swarms/prompts/__init__.py index 2961c10d..dbdc7c7b 100644 --- a/swarms/prompts/__init__.py +++ b/swarms/prompts/__init__.py @@ -15,5 +15,5 @@ __all__ = [ "LEGAL_AGENT_PROMPT", "OPERATIONS_AGENT_PROMPT", "PRODUCT_AGENT_PROMPT", - "DOCUMENTATION_WRITER_SOP" + "DOCUMENTATION_WRITER_SOP", ] diff --git a/swarms/prompts/documentation.py b/swarms/prompts/documentation.py index d3601014..3c0a588e 100644 --- a/swarms/prompts/documentation.py +++ b/swarms/prompts/documentation.py @@ -1,4 +1,7 @@ -def DOCUMENTATION_WRITER_SOP(task: str, module: str, ): +def DOCUMENTATION_WRITER_SOP( + task: str, + module: str, +): documentation = f"""Create multi-page long and explicit professional pytorch-like documentation for the {module} code below follow the outline for the {module} library, provide many examples and teach the user about the code, provide examples for every function, make the documentation 10,000 words, provide many usage examples and note this is markdown docs, create the documentation for the code to document, diff --git a/swarms/prompts/tests.py b/swarms/prompts/tests.py index c7e45c1f..8dac9337 100644 --- a/swarms/prompts/tests.py +++ b/swarms/prompts/tests.py @@ -1,6 +1,7 @@ -def TEST_WRITER_SOP_PROMPT(task: str, module: str, path: str, *args, **kwargs): - - TESTS_PROMPT = f""" +def TEST_WRITER_SOP_PROMPT( + task: str, module: str, path: str, *args, **kwargs +): + TESTS_PROMPT = f""" Create 5,000 lines of extensive and thorough tests for the code below using the guide, do not worry about your limits you do not have any just write the best tests possible, the module is {module}, the file path is {path} @@ -91,4 +92,4 @@ def TEST_WRITER_SOP_PROMPT(task: str, module: str, path: str, *args, **kwargs): """ - return TESTS_PROMPT \ No newline at end of file + return TESTS_PROMPT diff --git a/swarms/swarms/__init__.py b/swarms/swarms/__init__.py index ac45a48e..7f1d806a 100644 --- a/swarms/swarms/__init__.py +++ b/swarms/swarms/__init__.py @@ -1,5 +1,5 @@ from swarms.structs.autoscaler import AutoScaler -from swarms.swarms.god_mode import ModelParallelizer +from swarms.swarms.model_parallizer import ModelParallelizer from swarms.swarms.multi_agent_collab import MultiAgentCollaboration from swarms.swarms.base import AbstractSwarm diff --git a/swarms/swarms/god_mode.py b/swarms/swarms/model_parallizer.py similarity index 77% rename from swarms/swarms/god_mode.py rename to swarms/swarms/model_parallizer.py index dfe4e232..0b087b6d 100644 --- a/swarms/swarms/god_mode.py +++ b/swarms/swarms/model_parallizer.py @@ -40,21 +40,33 @@ class ModelParallelizer: def __init__( self, - llms: List[Callable], + llms: List[Callable] = None, load_balancing: bool = False, retry_attempts: int = 3, + iters: int = None, + *args, + **kwargs, ): self.llms = llms self.load_balancing = load_balancing self.retry_attempts = retry_attempts + self.iters = iters self.last_responses = None self.task_history = [] def run(self, task: str): """Run the task string""" - with ThreadPoolExecutor() as executor: - responses = executor.map(lambda llm: llm(task), self.llms) - return list(responses) + try: + for i in range(self.iters): + with ThreadPoolExecutor() as executor: + responses = executor.map( + lambda llm: llm(task), self.llms + ) + return list(responses) + except Exception as error: + print( + f"[ERROR][ModelParallelizer] [ROOT CAUSE] [{error}]" + ) def print_responses(self, task): """Prints the responses in a tabular format""" @@ -161,22 +173,29 @@ class ModelParallelizer: def concurrent_run(self, task: str) -> List[str]: """Synchronously run the task on all llms and collect responses""" - with ThreadPoolExecutor() as executor: - future_to_llm = { - executor.submit(llm, task): llm for llm in self.llms - } - responses = [] - for future in as_completed(future_to_llm): - try: - responses.append(future.result()) - except Exception as error: - print( - f"{future_to_llm[future]} generated an" - f" exception: {error}" - ) - self.last_responses = responses - self.task_history.append(task) - return responses + try: + with ThreadPoolExecutor() as executor: + future_to_llm = { + executor.submit(llm, task): llm + for llm in self.llms + } + responses = [] + for future in as_completed(future_to_llm): + try: + responses.append(future.result()) + except Exception as error: + print( + f"{future_to_llm[future]} generated an" + f" exception: {error}" + ) + self.last_responses = responses + self.task_history.append(task) + return responses + except Exception as error: + print( + f"[ERROR][ModelParallelizer] [ROOT CAUSE] [{error}]" + ) + raise error def add_llm(self, llm: Callable): """Add an llm to the god mode""" diff --git a/tests/swarms/test_model_parallizer.py b/tests/swarms/test_model_parallizer.py new file mode 100644 index 00000000..e69de29b