Feature/SK-944 + SK-934 | Add flag for deleting the virtual env

fedn/cli/client_cmd.py

@@ -6,8 +6,8 @@
 from fedn.common.exceptions import InvalidClientConfig
 from fedn.network.clients.client import Client
 
-from .main import main
-from .shared import CONTROLLER_DEFAULTS, apply_config, get_api_url, get_token, print_response
+from fedn.cli.main import main
+from fedn.cli.shared import CONTROLLER_DEFAULTS, apply_config, get_api_url, get_token, print_response
 
 
 def validate_client_config(config):

fedn/cli/controller_cmd.py

@@ -1,6 +1,6 @@
 import click
 
-from .main import main
+from fedn.cli.main import main
 
 
 @main.group("controller")

fedn/cli/run_cmd.py

@@ -10,10 +10,9 @@
 from fedn.network.combiner.combiner import Combiner
 from fedn.utils.dispatcher import Dispatcher, _read_yaml_file
 
-from .client_cmd import validate_client_config
-from .main import main
-from .shared import apply_config
-
+from fedn.cli.client_cmd import validate_client_config
+from fedn.cli.main import main
+from fedn.cli.shared import apply_config
 
 def get_statestore_config_from_file(init):
     """:param init:
@@ -36,7 +35,22 @@ def check_helper_config_file(config):
         exit(-1)
     return helper
 
-
+def check_yaml_exists(path):
+    """Check if fedn.yaml exists in the given path."""
+    yaml_file = os.path.join(path, "fedn.yaml")
+    if not os.path.exists(yaml_file):
+        logger.error(f"Could not find fedn.yaml in {path}")
+        click.echo(f"Could not find fedn.yaml in {path}")
+        exit(-1)
+    return yaml_file
+def delete_virtual_environment(keep_venv, dispatcher):
+    if not keep_venv:
+        # delete the virtualenv
+        if dispatcher.python_env_path:
+            logger.info(f"Removing virtualenv {dispatcher.python_env_path}")
+            shutil.rmtree(dispatcher.python_env_path)
+        else:
+            logger.warning("No virtualenv found to remove.")
 @main.group("run")
 @click.pass_context
 def run_cmd(ctx):
@@ -46,20 +60,18 @@ def run_cmd(ctx):
 @run_cmd.command("validate")
 @click.option("-p", "--path", required=True, help="Path to package directory containing fedn.yaml")
 @click.option("-i", "--input", required=True, help="Path to input model" )
-@click.option("-o", "--output", required=True,help="Path to write the output JSON containing validation metrics")
+@click.option("-o", "--output", required=True, help="Path to write the output JSON containing validation metrics")
+@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv")
 @click.pass_context
-def validate_cmd(ctx, path,input,output):
+def validate_cmd(ctx, path, input, output, keep_venv):
     """Execute 'validate' entrypoint in fedn.yaml.
 
     :param ctx:
     :param path: Path to folder containing fedn.yaml
     :type path: str
     """
     path = os.path.abspath(path)
-    yaml_file = os.path.join(path, "fedn.yaml")
-    if not os.path.exists(yaml_file):
-        logger.error(f"Could not find fedn.yaml in {path}")
-        exit(-1)
+    yaml_file = check_yaml_exists(path)
 
     config = _read_yaml_file(yaml_file)
     # Check that validate is defined in fedn.yaml under entry_points
@@ -70,28 +82,22 @@ def validate_cmd(ctx, path,input,output):
     dispatcher = Dispatcher(config, path)
     _ = dispatcher._get_or_create_python_env()
     dispatcher.run_cmd("validate {} {}".format(input, output))
-
-    # delete the virtualenv
-    if dispatcher.python_env_path:
-        logger.info(f"Removing virtualenv {dispatcher.python_env_path}")
-        shutil.rmtree(dispatcher.python_env_path)
+    delete_virtual_environment(keep_venv, dispatcher)
 @run_cmd.command("train")
 @click.option("-p", "--path", required=True, help="Path to package directory containing fedn.yaml")
 @click.option("-i", "--input", required=True, help="Path to input model parameters" )
-@click.option("-o", "--output", required=True,help="Path to write the updated model parameters ")
+@click.option("-o", "--output", required=True, help="Path to write the updated model parameters ")
+@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv")
 @click.pass_context
-def train_cmd(ctx, path,input,output):
+def train_cmd(ctx, path, input, output, keep_venv):
     """Execute 'train' entrypoint in fedn.yaml.
 
     :param ctx:
     :param path: Path to folder containing fedn.yaml
     :type path: str
     """
     path = os.path.abspath(path)
-    yaml_file = os.path.join(path, "fedn.yaml")
-    if not os.path.exists(yaml_file):
-        logger.error(f"Could not find fedn.yaml in {path}")
-        exit(-1)
+    yaml_file = check_yaml_exists(path)
 
     config = _read_yaml_file(yaml_file)
     # Check that train is defined in fedn.yaml under entry_points
@@ -102,57 +108,46 @@ def train_cmd(ctx, path,input,output):
     dispatcher = Dispatcher(config, path)
     _ = dispatcher._get_or_create_python_env()
     dispatcher.run_cmd("train {} {}".format(input, output))
+    delete_virtual_environment(keep_venv, dispatcher)
 
-    # delete the virtualenv
-    if dispatcher.python_env_path:
-        logger.info(f"Removing virtualenv {dispatcher.python_env_path}")
-        shutil.rmtree(dispatcher.python_env_path)
 @run_cmd.command("startup")
 @click.option("-p", "--path", required=True, help="Path to package directory containing fedn.yaml")
+@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv")
 @click.pass_context
-def startup_cmd(ctx, path):
+def startup_cmd(ctx, path, keep_venv):
     """Execute 'startup' entrypoint in fedn.yaml.
 
     :param ctx:
     :param path: Path to folder containing fedn.yaml
     :type path: str
     """
     path = os.path.abspath(path)
-    yaml_file = os.path.join(path, "fedn.yaml")
-    if not os.path.exists(yaml_file):
-        logger.error(f"Could not find fedn.yaml in {path}")
-        exit(-1)
+    yaml_file = check_yaml_exists(path)
 
     config = _read_yaml_file(yaml_file)
     # Check that startup is defined in fedn.yaml under entry_points
     if "startup" not in config["entry_points"]:
         logger.error("No startup command defined in fedn.yaml")
         exit(-1)
-
     dispatcher = Dispatcher(config, path)
     _ = dispatcher._get_or_create_python_env()
     dispatcher.run_cmd("startup")
+    delete_virtual_environment(keep_venv, dispatcher)
 
-    # delete the virtualenv
-    if dispatcher.python_env_path:
-        logger.info(f"Removing virtualenv {dispatcher.python_env_path}")
-        shutil.rmtree(dispatcher.python_env_path)
 
 @run_cmd.command("build")
 @click.option("-p", "--path", required=True, help="Path to package directory containing fedn.yaml")
+@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv")
 @click.pass_context
-def build_cmd(ctx, path):
+def build_cmd(ctx, path, keep_venv):
     """Execute 'build' entrypoint in fedn.yaml.
 
     :param ctx:
     :param path: Path to folder containing fedn.yaml
     :type path: str
     """
     path = os.path.abspath(path)
-    yaml_file = os.path.join(path, "fedn.yaml")
-    if not os.path.exists(yaml_file):
-        logger.error(f"Could not find fedn.yaml in {path}")
-        exit(-1)
+    yaml_file = check_yaml_exists(path)
 
     config = _read_yaml_file(yaml_file)
     # Check that build is defined in fedn.yaml under entry_points
@@ -163,11 +158,8 @@ def build_cmd(ctx, path):
     dispatcher = Dispatcher(config, path)
     _ = dispatcher._get_or_create_python_env()
     dispatcher.run_cmd("build")
+    delete_virtual_environment(keep_venv, dispatcher)
 
-    # delete the virtualenv
-    if dispatcher.python_env_path:
-        logger.info(f"Removing virtualenv {dispatcher.python_env_path}")
-        shutil.rmtree(dispatcher.python_env_path)
 
 
 @run_cmd.command("client")