Improve reset feature.

README.md

@@ -31,12 +31,14 @@ pip install -U autocode-py
 2. Prepare software to be processed as in the [`./example/client`](https://github.com/muazhari/autocode/tree/main/example/client) folder.
 3. Prepare deployment as in the [`./example/client/docker-compose.yml`](https://github.com/muazhari/autocode/blob/main/example/client/docker-compose.yml) file.
 4. Prepare controller as in the [`./example/controller.ipynb`](https://github.com/muazhari/autocode/blob/main/example/controller.ipynb) file.
-5. Instantiate `optimization` and execute `optimization.deploy()` in controller.
+5. Instantiate `optimization` then execute `optimization.deploy()` in controller.
 6. Open dashboard in `http://localhost:{dashboard_port}/` to see the process in real-time.
 7. Wait until all client are ready (need to wait for long time because the libraries need to be re-downloaded for each client).
 8. Execute `optimization.run()` in controller.
 9. Wait until the run is finished.
 10. Analyze and decide the best values.
+11. Execute `optimization.reset(keys=["clients"])` then `optimization.deploy()` to apply different client states.
+12. Try to execute `optimization.reset()` to totally reset the tool if needed (i.e. data inconsistency).
 
 ## Demo
 

autocode/__init__.py

@@ -2,7 +2,7 @@
 import sys
 from http import HTTPStatus
 from multiprocessing import Process
-from typing import Callable, Any, Dict, List
+from typing import Callable, Any, Dict, List, Literal, Optional
 
 import ray
 import uvicorn
@@ -110,14 +110,16 @@ def stop(self):
         self.server.terminate()
         ray.shutdown()
 
-    def reset(self):
+    def reset(self, keys: Optional[List[Literal["docker_clients", "clients", "objectives", "results"]]] = None):
         optimization_use_case: OptimizationUseCase = self.application_container.use_cases.optimization()
-        optimization_use_case.reset()
+        optimization_use_case.reset(
+            keys=keys
+        )
 
     def deploy(self, compose_files: List[str]) -> List[DockerClient]:
         optimization_use_case: OptimizationUseCase = self.application_container.use_cases.optimization()
         return optimization_use_case.deploy(
-            compose_files=compose_files
+            compose_files=compose_files,
         )
 
     def run(

autocode/use_case.py

@@ -3,12 +3,11 @@
 import re
 import uuid
 from multiprocessing import Lock
-from typing import Dict, List, Callable, Any, Coroutine, Optional
+from typing import Dict, List, Callable, Any, Coroutine, Optional, Literal
 
 import dill
 import numpy as np
 import ray
-import sqlalchemy.exc
 from langchain_core.output_parsers import PydanticToolsParser
 from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
 from langchain_core.runnables import RunnableSerializable
@@ -504,6 +503,7 @@ def deploy(self, compose_files: List[str]) -> List[DockerClient]:
         session.begin()
         session.exec(delete(Cache).where(Cache.key.startswith("results")))
         session.exec(delete(Cache).where(Cache.key == "objectives"))
+
         client_caches: List[Cache] = list(session.exec(select(Cache).where(Cache.key.startswith("clients"))).all())
         docker_clients: List[DockerClient] = []
 
@@ -571,19 +571,31 @@ def deploy(self, compose_files: List[str]) -> List[DockerClient]:
 
         return docker_clients
 
-    def reset(self):
+    def reset(self, keys: Optional[List[Literal["docker_clients", "clients", "objectives", "results"]]] = None):
         session: Session = self.one_datastore.get_session()
         session.begin()
-        try:
-            docker_client_caches = list(session.exec(select(Cache).where(Cache.key.startswith("docker_clients"))).all())
-            for docker_client_cache in docker_client_caches:
-                docker_client: DockerClient = dill.loads(docker_client_cache.value)
-                docker_client.compose.down()
-        except sqlalchemy.exc.OperationalError:
-            pass
+        if keys is None:
+            SQLModel.metadata.drop_all(self.one_datastore.engine)
+            SQLModel.metadata.create_all(self.one_datastore.engine)
+        else:
+            if "docker_clients" in keys:
+                docker_client_caches = list(
+                    session.exec(select(Cache).where(Cache.key.startswith("docker_clients"))).all()
+                )
+                for docker_client_cache in docker_client_caches:
+                    docker_client: DockerClient = dill.loads(docker_client_cache.value)
+                    docker_client.compose.down()
+                    session.delete(docker_client_cache)
+
+            if "clients" in keys:
+                session.exec(delete(Cache).where(Cache.key.startswith("clients")))
+
+            if "objectives" in keys:
+                session.exec(delete(Cache).where(Cache.key == "objectives"))
+
+            if "results" in keys:
+                session.exec(delete(Cache).where(Cache.key.startswith("results")))
 
-        SQLModel.metadata.drop_all(self.one_datastore.engine)
-        SQLModel.metadata.create_all(self.one_datastore.engine)
         session.commit()
         session.close()