diff --git a/aiopslab-applications b/aiopslab-applications index 5221ef96..48e03edb 160000 --- a/aiopslab-applications +++ b/aiopslab-applications @@ -1 +1 @@ -Subproject commit 5221ef962879546bb3c977c8a256ce117697b9f2 +Subproject commit 48e03edb4732468331b6963bc4644e8bae08fac1 diff --git a/aiopslab/generators/fault/inject_virtual.py b/aiopslab/generators/fault/inject_virtual.py index 8b6680d0..2ee4089d 100644 --- a/aiopslab/generators/fault/inject_virtual.py +++ b/aiopslab/generators/fault/inject_virtual.py @@ -8,6 +8,7 @@ from aiopslab.service.kubectl import KubeCtl from aiopslab.service.helm import Helm +from aiopslab.service.dock import Docker from aiopslab.generators.fault.base import FaultInjector from aiopslab.service.apps.base import Application from aiopslab.paths import TARGET_MICROSERVICES @@ -18,6 +19,7 @@ def __init__(self, namespace: str): super().__init__(namespace) self.namespace = namespace self.kubectl = KubeCtl() + self.docker = Docker() self.mongo_service_pod_map = { "url-shorten-mongodb": "url-shorten-service", } @@ -248,7 +250,35 @@ def recover_wrong_bin_usage(self, microservices: list[str]): self.kubectl.exec_command(apply_command) print(f"Recovered from wrong binary usage fault for service: {service}") - + + def inject_container_stop(self, microservices: list[str]): + """Inject a fault to stop a container.""" + for service in microservices: + self.docker.get_container(service).stop() + print(f"Stopped container {service}.") + + print("Waiting for faults to propagate...") + time.sleep(15) + print("Faults propagated.") + + def recover_container_stop(self, microservices: list[str]): + for service in microservices: + self.docker.get_container(service).start() + print(f"Started container {service}.") + + def inject_model_misconfig(self, microservices: list[str]): + """Inject a fault to misconfigure the model in the Flower application.""" + for service in microservices: + command = f""" docker exec -it {service} sh -c "sed -i '24s/84/80/' /app/.flwr/apps/*/task.py" """ + self.docker.exec_command(command) + print(f"Changed model configuration for service: {service}") + + def recover_model_misconfig(self, microservices: list[str]): + for service in microservices: + command = f""" docker exec -it {service} sh -c "sed -i '24s/80/84/' /app/.flwr/apps/*/task.py" """ + self.docker.exec_command(command) + print(f"Recovered model configuration for service: {service}") + ############# HELPER FUNCTIONS ################ def _wait_for_pods_ready(self, microservices: list[str], timeout: int = 30): for service in microservices: diff --git a/aiopslab/orchestrator/actions/base.py b/aiopslab/orchestrator/actions/base.py index 1baf32a2..2aa9d49f 100644 --- a/aiopslab/orchestrator/actions/base.py +++ b/aiopslab/orchestrator/actions/base.py @@ -8,6 +8,7 @@ from datetime import datetime, timedelta from aiopslab.utils.actions import action, read, write from aiopslab.service.kubectl import KubeCtl +from aiopslab.service.dock import Docker from aiopslab.service.shell import Shell # from aiopslab.observer import initialize_pod_and_service_lists @@ -22,7 +23,7 @@ class TaskActions: @read def get_logs(namespace: str, service: str) -> str: """ - Collects relevant log data from a pod using Kubectl. + Collects relevant log data from a pod using Kubectl or from a container with Docker. Args: namespace (str): The namespace in which the service is running. @@ -31,26 +32,35 @@ def get_logs(namespace: str, service: str) -> str: Returns: str | dict | list[dicts]: Log data as a structured object or a string. """ - kubectl = KubeCtl() - try: - if namespace == "test-social-network": - user_service_pod = kubectl.get_pod_name(namespace, f"app={service}") - elif namespace == "test-hotel-reservation": - user_service_pod = kubectl.get_pod_name( - namespace, f"io.kompose.service={service}" - ) - elif namespace == "astronomy-shop": - user_service_pod = kubectl.get_pod_name( - namespace, f"app.kubernetes.io/name={service}" - ) - elif namespace == "default" and "wrk2-job" in service: - user_service_pod = kubectl.get_pod_name(namespace, f"job-name=wrk2-job") - else: - raise Exception - logs = kubectl.get_pod_logs(user_service_pod, namespace) - except Exception as e: - return "Error: Your service/namespace does not exist. Use kubectl to check." - + if namespace == "docker": + docker = Docker() + try: + logs = docker.get_logs(service) + except Exception as e: + return "Error: Your service does not exist. Use docker to check." + + else: + kubectl = KubeCtl() + try: + if namespace == "test-social-network": + user_service_pod = kubectl.get_pod_name(namespace, f"app={service}") + elif namespace == "test-hotel-reservation": + user_service_pod = kubectl.get_pod_name( + namespace, f"io.kompose.service={service}" + ) + elif namespace == "astronomy-shop": + user_service_pod = kubectl.get_pod_name( + namespace, f"app.kubernetes.io/name={service}" + ) + elif namespace == "default" and "wrk2-job" in service: + user_service_pod = kubectl.get_pod_name(namespace, f"job-name=wrk2-job") + else: + raise Exception + logs = kubectl.get_pod_logs(user_service_pod, namespace) + except Exception as e: + return "Error: Your service/namespace does not exist. Use kubectl to check." + + print(logs) logs = "\n".join(logs.split("\n")) return logs @@ -71,6 +81,9 @@ def exec_shell(command: str) -> str: """ if "kubectl edit" in command or "edit svc" in command: return "Error: Cannot use `kubectl edit`. Use `kubectl patch` instead." + + if "docker logs -f" in command: + return "Error: Cannot use `docker logs -f`. Use `docker logs` instead." return Shell.exec(command) diff --git a/aiopslab/orchestrator/orchestrator.py b/aiopslab/orchestrator/orchestrator.py index d3fbc455..cdf109f1 100644 --- a/aiopslab/orchestrator/orchestrator.py +++ b/aiopslab/orchestrator/orchestrator.py @@ -45,24 +45,26 @@ def init_problem(self, problem_id: str): self.session = Session() print(f"Session ID: {self.session.session_id}") prob = self.probs.get_problem_instance(problem_id) + deployment = self.probs.get_problem_deployment(problem_id) self.session.set_problem(prob, pid=problem_id) self.session.set_agent(self.agent_name) - print("Setting up OpenEBS...") + if deployment != "docker": + print("Setting up OpenEBS...") - # Install OpenEBS - self.kubectl.exec_command( - "kubectl apply -f https://openebs.github.io/charts/openebs-operator.yaml" - ) - self.kubectl.exec_command( - "kubectl patch storageclass openebs-hostpath -p '{\"metadata\": {\"annotations\":{\"storageclass.kubernetes.io/is-default-class\":\"true\"}}}'" - ) - self.kubectl.wait_for_ready("openebs") - print("OpenEBS setup completed.") + # Install OpenEBS + self.kubectl.exec_command( + "kubectl apply -f https://openebs.github.io/charts/openebs-operator.yaml" + ) + self.kubectl.exec_command( + "kubectl patch storageclass openebs-hostpath -p '{\"metadata\": {\"annotations\":{\"storageclass.kubernetes.io/is-default-class\":\"true\"}}}'" + ) + self.kubectl.wait_for_ready("openebs") + print("OpenEBS setup completed.") - # Setup and deploy Prometheus - self.prometheus = Prometheus() - self.prometheus.deploy() + # Setup and deploy Prometheus + self.prometheus = Prometheus() + self.prometheus.deploy() # deploy service prob.app.delete() @@ -200,13 +202,15 @@ async def start_problem(self, max_steps: int): # But this will take more time. # if not self.session.problem.sys_status_after_recovery(): self.session.problem.app.cleanup() - self.prometheus.teardown() - print("Uninstalling OpenEBS...") - self.kubectl.exec_command("kubectl delete sc openebs-hostpath openebs-device --ignore-not-found") - self.kubectl.exec_command( - "kubectl delete -f https://openebs.github.io/charts/openebs-operator.yaml" - ) - self.kubectl.wait_for_namespace_deletion("openebs") + + if self.session.problem.namespace != "docker": + self.prometheus.teardown() + print("Uninstalling OpenEBS...") + self.kubectl.exec_command("kubectl delete sc openebs-hostpath openebs-device --ignore-not-found") + self.kubectl.exec_command( + "kubectl delete -f https://openebs.github.io/charts/openebs-operator.yaml" + ) + self.kubectl.wait_for_namespace_deletion("openebs") self.execution_end_time = time.time() total_execution_time = self.execution_end_time - self.execution_start_time diff --git a/aiopslab/orchestrator/problems/flower_model_misconfig/__init__.py b/aiopslab/orchestrator/problems/flower_model_misconfig/__init__.py new file mode 100644 index 00000000..1db81e08 --- /dev/null +++ b/aiopslab/orchestrator/problems/flower_model_misconfig/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from .model_misconfig import ( + FlowerModelMisconfigDetection +) \ No newline at end of file diff --git a/aiopslab/orchestrator/problems/flower_model_misconfig/model_misconfig.py b/aiopslab/orchestrator/problems/flower_model_misconfig/model_misconfig.py new file mode 100644 index 00000000..9495d716 --- /dev/null +++ b/aiopslab/orchestrator/problems/flower_model_misconfig/model_misconfig.py @@ -0,0 +1,96 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Model misconfiguration fault in the Flower application.""" + +import time +from typing import Any + +from aiopslab.orchestrator.tasks import * +from aiopslab.service.dock import Docker +from aiopslab.service.apps.flower import Flower +from aiopslab.paths import TARGET_MICROSERVICES +from aiopslab.session import SessionItem +from aiopslab.generators.fault.inject_virtual import VirtualizationFaultInjector + + +class FlowerModelMisconfigBaseTask: + def __init__(self, faulty_service: str = "user-service"): + self.app = Flower() + self.docker = Docker() + self.namespace = self.app.namespace + self.faulty_service = faulty_service + self.train_dir = TARGET_MICROSERVICES / "flower" + + def start_workload(self): + print("== Start Workload ==") + command = "flwr run train local-deployment" + self.docker.exec_command(command, cwd=self.train_dir) + + path = "/app/.flwr/apps" + check = f""" docker exec -it {self.faulty_service} sh -c "test -d {path} && echo 'exists'" """ + + print("Waiting for workload to start...") + while True: + exists = self.docker.exec_command(check) + if exists.strip() == "exists": + break + time.sleep(1) + print("Workload started successfully.") + + # Inject fault after workload starts, since the required files are created during the workload + print("Injecting fault...") + self.inject_fault(inject=True) + + print("Waiting for faults to propagate...") + while True: + logs = self.docker.get_logs(self.faulty_service) + if "error" in logs.lower(): + break + time.sleep(1) + print("Faults propagated.") + + def inject_fault(self, inject: bool = False): + print("== Fault Injection ==") + if inject: + injector = VirtualizationFaultInjector(namespace=self.namespace) + injector._inject( + fault_type="model_misconfig", + microservices=[self.faulty_service], + ) + print(f"Service: {self.faulty_service} | Namespace: {self.namespace}\n") + else: + print("Fault injection skipped.") + + def recover_fault(self): + print("== Fault Recovery ==") + injector = VirtualizationFaultInjector(namespace=self.namespace) + injector._recover( + fault_type="model_misconfig", + microservices=[self.faulty_service], + ) + print(f"Service: {self.faulty_service} | Namespace: {self.namespace}\n") + + +################## Detection Problem ################## +class FlowerModelMisconfigDetection(FlowerModelMisconfigBaseTask, DetectionTask): + def __init__(self, faulty_service: str = "clientapp-1"): + FlowerModelMisconfigBaseTask.__init__(self, faulty_service=faulty_service) + DetectionTask.__init__(self, self.app) + + def eval(self, soln: Any, trace: list[SessionItem], duration: float): + print("== Evaluation ==") + expected_solution = "Yes" + + if isinstance(soln, str): + if soln.strip().lower() == expected_solution.lower(): + print(f"Correct detection: {soln}") + self.add_result("Detection Accuracy", "Correct") + else: + print(f"Incorrect detection: {soln}") + self.add_result("Detection Accuracy", "Incorrect") + else: + print("Invalid solution format") + self.add_result("Detection Accuracy", "Invalid Format") + + return super().eval(soln, trace, duration) diff --git a/aiopslab/orchestrator/problems/flower_node_stop/__init__.py b/aiopslab/orchestrator/problems/flower_node_stop/__init__.py new file mode 100644 index 00000000..fb33063c --- /dev/null +++ b/aiopslab/orchestrator/problems/flower_node_stop/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from .node_stop import ( + FlowerNodeStopDetection +) \ No newline at end of file diff --git a/aiopslab/orchestrator/problems/flower_node_stop/node_stop.py b/aiopslab/orchestrator/problems/flower_node_stop/node_stop.py new file mode 100644 index 00000000..79596976 --- /dev/null +++ b/aiopslab/orchestrator/problems/flower_node_stop/node_stop.py @@ -0,0 +1,69 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Docker node stop fault problem in the Flower application.""" + +from typing import Any + +from aiopslab.orchestrator.tasks import * +from aiopslab.service.dock import Docker +from aiopslab.service.apps.flower import Flower +from aiopslab.paths import TARGET_MICROSERVICES +from aiopslab.session import SessionItem +from aiopslab.generators.fault.inject_virtual import VirtualizationFaultInjector + + +class FlowerNodeStopBaseTask: + def __init__(self, faulty_service: str = "user-service"): + self.app = Flower() + self.docker = Docker() + self.namespace = self.app.namespace + self.faulty_service = faulty_service + self.train_dir = TARGET_MICROSERVICES / "flower" + + def start_workload(self): + print("== Start Workload ==") + command = "flwr run train local-deployment" + self.docker.exec_command(command, cwd=self.train_dir) + + def inject_fault(self): + print("== Fault Injection ==") + injector = VirtualizationFaultInjector(namespace=self.namespace) + injector._inject( + fault_type="container_stop", + microservices=[self.faulty_service], + ) + print(f"Service: {self.faulty_service} | Namespace: {self.namespace}\n") + + def recover_fault(self): + print("== Fault Recovery ==") + injector = VirtualizationFaultInjector(namespace=self.namespace) + injector._recover( + fault_type="container_stop", + microservices=[self.faulty_service], + ) + print(f"Service: {self.faulty_service} | Namespace: {self.namespace}\n") + + +################## Detection Problem ################## +class FlowerNodeStopDetection(FlowerNodeStopBaseTask, DetectionTask): + def __init__(self, faulty_service: str = "supernode-1"): + FlowerNodeStopBaseTask.__init__(self, faulty_service=faulty_service) + DetectionTask.__init__(self, self.app) + + def eval(self, soln: Any, trace: list[SessionItem], duration: float): + print("== Evaluation ==") + expected_solution = "Yes" + + if isinstance(soln, str): + if soln.strip().lower() == expected_solution.lower(): + print(f"Correct detection: {soln}") + self.add_result("Detection Accuracy", "Correct") + else: + print(f"Incorrect detection: {soln}") + self.add_result("Detection Accuracy", "Incorrect") + else: + print("Invalid solution format") + self.add_result("Detection Accuracy", "Invalid Format") + + return super().eval(soln, trace, duration) diff --git a/aiopslab/orchestrator/problems/registry.py b/aiopslab/orchestrator/problems/registry.py index 7a0bcae8..9a913227 100644 --- a/aiopslab/orchestrator/problems/registry.py +++ b/aiopslab/orchestrator/problems/registry.py @@ -27,6 +27,8 @@ from aiopslab.orchestrator.problems.redeploy_without_pv import * from aiopslab.orchestrator.problems.wrong_bin_usage import * from aiopslab.orchestrator.problems.operator_misoperation import * +from aiopslab.orchestrator.problems.flower_node_stop import * +from aiopslab.orchestrator.problems.flower_model_misconfig import * class ProblemRegistry: @@ -211,7 +213,14 @@ def __init__(self): # "operator_security_context_fault-localization-1": K8SOperatorSecurityContextFaultLocalization, # "operator_wrong_update_strategy-detection-1": K8SOperatorWrongUpdateStrategyDetection, # "operator_wrong_update_strategy-localization-1": K8SOperatorWrongUpdateStrategyLocalization, + # Flower + "flower_node_stop-detection": FlowerNodeStopDetection, + "flower_model_misconfig-detection": FlowerModelMisconfigDetection, } + self.DOCKER_REGISTRY = [ + "flower_node_stop-detection", + "flower_model_misconfig-detection", + ] def get_problem_instance(self, problem_id: str): if problem_id not in self.PROBLEM_REGISTRY: @@ -231,3 +240,8 @@ def get_problem_count(self, task_type: str = None): if task_type: return len([k for k in self.PROBLEM_REGISTRY.keys() if task_type in k]) return len(self.PROBLEM_REGISTRY) + + def get_problem_deployment(self, problem_id: str): + if problem_id in self.DOCKER_REGISTRY: + return "docker" + return "k8s" diff --git a/aiopslab/paths.py b/aiopslab/paths.py index 6aad9b2a..2391bc41 100644 --- a/aiopslab/paths.py +++ b/aiopslab/paths.py @@ -34,3 +34,4 @@ ASTRONOMY_SHOP_METADATA = BASE_DIR / "service" / "metadata" / "astronomy-shop.json" TIDB_METADATA = BASE_DIR / "service" / "metadata" / "tidb-with-operator.json" FLIGHT_TICKET_METADATA = BASE_DIR / "service" / "metadata" / "flight-ticket.json" +FLOWER_METADATA = BASE_DIR / "service" / "metadata" / "flower.json" \ No newline at end of file diff --git a/aiopslab/service/apps/base.py b/aiopslab/service/apps/base.py index 02a30b15..b88a0ee6 100644 --- a/aiopslab/service/apps/base.py +++ b/aiopslab/service/apps/base.py @@ -15,6 +15,7 @@ def __init__(self, config_file: str): self.helm_deploy = True self.helm_configs = {} self.k8s_deploy_path = None + self.docker_deploy_path = None def load_app_json(self): """Load (basic) application metadata into attributes. @@ -35,6 +36,9 @@ def load_app_json(self): if "K8S Deploy Path" in metadata: self.k8s_deploy_path = TARGET_MICROSERVICES / metadata["K8S Deploy Path"] + + if "Docker Deploy Path" in metadata: + self.docker_deploy_path = TARGET_MICROSERVICES / metadata["Docker Deploy Path"] def get_app_json(self) -> dict: """Get application metadata in JSON format. diff --git a/aiopslab/service/apps/flower.py b/aiopslab/service/apps/flower.py new file mode 100644 index 00000000..22834d3e --- /dev/null +++ b/aiopslab/service/apps/flower.py @@ -0,0 +1,29 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from aiopslab.service.dock import Docker +from aiopslab.service.apps.base import Application +from aiopslab.paths import FLOWER_METADATA + + +class Flower(Application): + def __init__(self): + super().__init__(FLOWER_METADATA) + self.docker = Docker() + + self.load_app_json() + + def deploy(self): + """Deploy the docker compose file.""" + print("Deploying docker compose files") + self.docker.compose_up(self.docker_deploy_path) + + def delete(self): + """Stop the docker containers.""" + print("Stopping the docker containers") + self.docker.compose_down(self.docker_deploy_path) + + def cleanup(self): + """Delete all stopped docker containers.""" + print("Deleting stopped containers") + self.docker.cleanup() diff --git a/aiopslab/service/dock.py b/aiopslab/service/dock.py new file mode 100644 index 00000000..149ff1cc --- /dev/null +++ b/aiopslab/service/dock.py @@ -0,0 +1,57 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Interface to Docker controller service.""" + +import docker +import subprocess + + +class Docker: + def __init__(self): + self.client = docker.from_env() + + def list_containers(self): + """Get all containers.""" + return self.client.containers.list() + + def get_container(self, container_id): + """Get a container by ID.""" + return self.client.containers.get(container_id) + + def get_logs(self, container_id): + """Get logs for a container.""" + return self.get_container(container_id).logs().decode("utf-8") + + def compose_up(self, cwd): + """Run docker-compose up.""" + command = "docker compose up -d" + return self.exec_command(command, cwd=cwd) + + def compose_down(self, cwd): + """Run docker-compose down.""" + command = "docker compose down" + return self.exec_command(command, cwd=cwd) + + def cleanup(self): + """Remove the stopped docker containers.""" + command = "docker container prune -f" + return self.exec_command(command) + + def exec_command(self, command: str, input_data=None, cwd=None): + """Execute an arbitrary command.""" + if input_data is not None: + input_data = input_data.encode("utf-8") + try: + out = subprocess.run( + command, + input=input_data, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + shell=True, + cwd=cwd, + ) + if out is not None: + return out.stdout.decode("utf-8") + except subprocess.CalledProcessError as e: + return e.stderr.decode("utf-8") diff --git a/aiopslab/service/metadata/flower.json b/aiopslab/service/metadata/flower.json new file mode 100644 index 00000000..87653228 --- /dev/null +++ b/aiopslab/service/metadata/flower.json @@ -0,0 +1,10 @@ +{ + "Name": "Flower", + "Namespace": "docker", + "Desc": "A federated learning application to train models on edge devices and aggregate them on a central server. It consists of four main components, the serverapp which is the central server, the superlink which is the communication link for the serverapp, the clientapp which is the edge device application, and the supernode which is the communication link for the clientapp. The serverapp and superlink are used to aggregate models from the clientapp and supernode, while the clientapp and supernode are used to train models on edge devices. The current deployment consists of one server and two clients.", + "Supported Operations": [ + "Train models on edge devices", + "Aggregate models on a central server" + ], + "Docker Deploy Path": "flower" +} \ No newline at end of file diff --git a/aiopslab/service/shell.py b/aiopslab/service/shell.py index 815589a3..af482bac 100644 --- a/aiopslab/service/shell.py +++ b/aiopslab/service/shell.py @@ -20,16 +20,16 @@ class Shell: def exec(command: str, input_data=None, cwd=None): """Execute a shell command on localhost, via SSH, or inside kind's control-plane container.""" k8s_host = config.get("k8s_host", "localhost") # Default to localhost - + if k8s_host == "kind": return Shell.docker_exec("kind-control-plane", command) elif k8s_host == "localhost": - print( - "[WARNING] Running commands on localhost is not recommended. " - "This may pose safety and security risks when using an AI agent locally. " - "I hope you know what you're doing!!!" - ) + # print( + # "[WARNING] Running commands on localhost is not recommended. " + # "This may pose safety and security risks when using an AI agent locally. " + # "I hope you know what you're doing!!!" + # ) return Shell.local_exec(command, input_data, cwd) else: @@ -50,14 +50,15 @@ def local_exec(command: str, input_data=None, cwd=None): stderr=subprocess.PIPE, shell=True, cwd=cwd, + timeout=10, # need to account for this properly ) - if out.stderr or out.returncode != 0: + if out.returncode != 0: error_message = out.stderr.decode("utf-8") print(f"[ERROR] Command execution failed: {error_message}") return error_message else: - output_message = out.stdout.decode("utf-8") + output_message = out.stdout.decode("utf-8") + out.stderr.decode("utf-8") print(output_message) return output_message diff --git a/clients/llama.py b/clients/llama.py new file mode 100644 index 00000000..a05795e4 --- /dev/null +++ b/clients/llama.py @@ -0,0 +1,66 @@ +""" +Naive LLaMA client (with shell access) for AIOpsLab. +""" + +import asyncio + +from aiopslab.orchestrator import Orchestrator +from clients.utils.llm import LLaMAClient +from clients.utils.templates import DOCS + + +class Agent: + def __init__(self): + self.history = [] + self.llm = LLaMAClient() + + def init_context(self, problem_desc: str, instructions: str, apis: str): + """Initialize the context for the agent.""" + + self.telemetry_apis = self._filter_dict(apis, lambda k, _: "get_logs" in k) + self.shell_api = self._filter_dict(apis, lambda k, _: "exec_shell" in k) + self.submit_api = self._filter_dict(apis, lambda k, _: "submit" in k) + stringify_apis = lambda apis: "\n\n".join( + [f"{k}\n{v}" for k, v in apis.items()] + ) + + self.system_message = DOCS.format( + prob_desc=problem_desc, + telemetry_apis=stringify_apis(self.telemetry_apis), + shell_api=stringify_apis(self.shell_api), + submit_api=stringify_apis(self.submit_api), + ) + + self.task_message = instructions + + self.history.append({"role": "system", "content": self.system_message}) + self.history.append({"role": "user", "content": self.task_message}) + + async def get_action(self, input) -> str: + """Wrapper to interface the agent with OpsBench. + + Args: + input (str): The input from the orchestrator/environment. + + Returns: + str: The response from the agent. + """ + self.history.append({"role": "user", "content": input}) + response = self.llm.run(self.history) + self.history.append({"role": "assistant", "content": response[0]}) + return response[0] + + def _filter_dict(self, dictionary, filter_func): + return {k: v for k, v in dictionary.items() if filter_func(k, v)} + + +if __name__ == "__main__": + agent = Agent() + + orchestrator = Orchestrator() + orchestrator.register_agent(agent, name="llama-w-shell") + + pid = "flower_model_misconfig-detection" + problem_desc, instructs, apis = orchestrator.init_problem(pid) + agent.init_context(problem_desc, instructs, apis) + asyncio.run(orchestrator.start_problem(max_steps=10)) diff --git a/clients/utils/llm.py b/clients/utils/llm.py index e473524b..5ecaebf2 100644 --- a/clients/utils/llm.py +++ b/clients/utils/llm.py @@ -6,6 +6,7 @@ import os from openai import OpenAI +from groq import Groq from pathlib import Path import json from dotenv import load_dotenv @@ -229,3 +230,43 @@ def run(self, payload: list[dict[str, str]]) -> list[str]: self.cache.add_to_cache(payload, response) self.cache.save_cache() return response + + +class LLaMAClient: + """Abstraction for Meta's LLaMA-3 model.""" + + def __init__(self): + self.cache = Cache() + + def inference(self, payload: list[dict[str, str]]) -> list[str]: + if self.cache is not None: + cache_result = self.cache.get_from_cache(payload) + if cache_result is not None: + return cache_result + + client = Groq(api_key=os.getenv("GROQ_API_KEY")) + try: + response = client.chat.completions.create( + messages=payload, + model="llama-3.1-8b-instant", + max_tokens=1024, + temperature=0.5, + top_p=0.95, + frequency_penalty=0.0, + presence_penalty=0.0, + n=1, + timeout=60, + stop=[], + ) + except Exception as e: + print(f"Exception: {repr(e)}") + raise e + + return [c.message.content for c in response.choices] # type: ignore + + def run(self, payload: list[dict[str, str]]) -> list[str]: + response = self.inference(payload) + if self.cache is not None: + self.cache.add_to_cache(payload, response) + self.cache.save_cache() + return response diff --git a/pyproject.toml b/pyproject.toml index b97b2f0f..eeb5f7e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,8 @@ python-dotenv = "^1.0.1" vllm = "^0.7.3" transformers = "^4.49.0" fastapi = "^0.115.12" +groq = "^0.28.0" +flwr = "^1.19.0" [build-system]