Remove containers and networks in pytest_sessionfinish

Signed-off-by: noopur <[email protected]>

tests/end_to_end/conftest.py

@@ -12,6 +12,7 @@
 from tests.end_to_end.utils.logger import configure_logging
 from tests.end_to_end.utils.logger import logger as log
 from tests.end_to_end.utils.conftest_helper import parse_arguments
+import tests.end_to_end.utils.docker_helper as dh
 
 
 def pytest_addoption(parser):
@@ -191,6 +192,11 @@ def pytest_sessionfinish(session, exitstatus):
     if os.path.exists(cache_dir):
         shutil.rmtree(cache_dir, ignore_errors=False)
         log.debug(f"Cleared .pytest_cache directory at {cache_dir}")
+
+    # Cleanup docker containers related to aggregator and collaborators, if any.
+    dh.cleanup_docker_containers(list_of_containers=["aggregator", "collaborator*"])
+    # Cleanup docker network created for openfl, if any.
+    dh.remove_docker_network(["openfl"])
 
 
 def pytest_configure(config):

tests/end_to_end/utils/docker_helper.py

@@ -12,35 +12,40 @@
 log = logging.getLogger(__name__)
 
 
-def remove_docker_network():
+def remove_docker_network(list_of_networks=[constants.DOCKER_NETWORK_NAME]):
     """
     Remove docker network.
+    Args:
+        list_of_networks (list): List of network names to remove.
     """
     client = get_docker_client()
-    networks = client.networks.list(names=[constants.DOCKER_NETWORK_NAME])
+    networks = client.networks.list(names=list_of_networks)
     if not networks:
-        log.debug(f"Network {constants.DOCKER_NETWORK_NAME} does not exist")
+        log.debug(f"Network(s) {list_of_networks} does not exist")
         return
 
     for network in networks:
         log.debug(f"Removing network: {network.name}")
         network.remove()
-    log.debug("Docker network removed successfully")
+    log.debug(f"Docker network(s) {list_of_networks} removed successfully")
 
 
-def create_docker_network():
+def create_docker_network(list_of_networks=[constants.DOCKER_NETWORK_NAME]):
     """
     Create docker network.
+    Args:
+        list_of_networks (list): List of network names to create.
     """
     client = get_docker_client()
-    networks = client.networks.list(names=[constants.DOCKER_NETWORK_NAME])
+    networks = client.networks.list(names=list_of_networks)
     if networks:
-        log.info(f"Network {constants.DOCKER_NETWORK_NAME} already exists")
+        log.info(f"Network(s) {list_of_networks} already exists")
         return
 
-    log.debug(f"Creating network: {constants.DOCKER_NETWORK_NAME}")
-    network = client.networks.create(constants.DOCKER_NETWORK_NAME)
-    log.info(f"Network {network.name} created successfully")
+    for network_name in list_of_networks:
+        log.debug(f"Creating network: {network_name}")
+        _ = client.networks.create(network_name)
+    log.info(f"Docker network(s) {list_of_networks} created successfully")
 
 
 def check_docker_image():
@@ -143,24 +148,24 @@ def get_docker_client():
     return client
 
 
-def cleanup_docker_containers():
+def cleanup_docker_containers(list_of_containers=["aggregator", "collaborator*"]):
     """
     Cleanup the docker containers meant for openfl.
+    Args:
+        list_of_containers: List of container names to cleanup.
     """
     log.debug("Cleaning up docker containers")
 
     client = get_docker_client()
 
-    # List all containers related to openfl
-    agg_containers = client.containers.list(all=True, filters={"name": "aggregator"})
-    col_containers = client.containers.list(all=True, filters={"name": "collaborator*"})
-    containers = agg_containers + col_containers
-    container_names = []
-    # Stop and remove all containers
-    for container in containers:
-        container.stop()
-        container.remove()
-        container_names.append(container.name)
-
-    if containers:
-        log.info(f"Docker containers {container_names} cleaned up successfully")
+    for container_name in list_of_containers:
+        containers = client.containers.list(all=True, filters={"name": container_name})
+        container_names = []
+        # Stop and remove all containers
+        for container in containers:
+            container.stop()
+            container.remove()
+            container_names.append(container.name)
+
+        if containers:
+            log.info(f"Docker containers {container_names} cleaned up successfully")

tests/end_to_end/utils/federation_helper.py

@@ -689,7 +689,7 @@ def setup_collaborator_data(collaborators, model_name, local_bind_path):
         model_name (str): Model name
         local_bind_path (str): Local bind path
     """
-    if not pre_existing_data(collaborators, model_name, local_bind_path):
+    if not pre_existing_data(collaborators):
         download_data(collaborators, model_name, local_bind_path)
 
     log.info("Data setup is complete for all the collaborators")
@@ -709,19 +709,12 @@ def download_data(collaborators, model_name, local_bind_path):
         raise ex.DataSetupException(f"Failed to copy data setup file: {e}")
 
     log.info("Downloading the data for the model. This will take some time to complete based on the data size ..")
-    error_msg = f"Failed to download data for {model_name}"
     try:
-        return_code, _, error = run_command(
-            f"python -v {constants.DATA_SETUP_FILE} {len(collaborators)}",
-            workspace_path=local_bind_path,
-            error_msg=error_msg,
-            return_error=True,
-        )
-        if return_code !=0 or error:
-            raise ex.DataSetupException(f"{error_msg}: {error}")
-
+        # Data setup file has logic to show progress bar
+        # Using os.system to make sure it is displayed
+        os.system(f"cd {local_bind_path}; python {constants.DATA_SETUP_FILE} {len(collaborators)}")
     except Exception:
-        raise ex.DataSetupException(f"Failed to download data for {model_name}") # Do not print error, it maybe too long
+        raise ex.DataSetupException(f"Failed to download data for {model_name}")
 
     try:
         # Move the data to the respective workspace based on the index