From e67db3021be96de33d31d3fc428c766703b953e8 Mon Sep 17 00:00:00 2001 From: sowmyasri Date: Mon, 19 Aug 2024 11:10:38 +0200 Subject: [PATCH] code clean up --- fedn/cli/run_cmd.py | 37 ++++++++++++++++++------------------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/fedn/cli/run_cmd.py b/fedn/cli/run_cmd.py index 93323b575..fb6179b7c 100644 --- a/fedn/cli/run_cmd.py +++ b/fedn/cli/run_cmd.py @@ -43,14 +43,12 @@ def check_yaml_exists(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.") +def delete_virtual_environment(dispatcher): + 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): @@ -61,7 +59,7 @@ def run_cmd(ctx): @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("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv") +@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set doesn't remove venv") @click.pass_context def validate_cmd(ctx, path, input, output, keep_venv): """Execute 'validate' entrypoint in fedn.yaml. @@ -82,12 +80,13 @@ def validate_cmd(ctx, path, input, output, keep_venv): dispatcher = Dispatcher(config, path) _ = dispatcher._get_or_create_python_env() dispatcher.run_cmd("validate {} {}".format(input, output)) - delete_virtual_environment(keep_venv, dispatcher) + if not keep_venv: + delete_virtual_environment(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("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv") +@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set doesn't remove venv") @click.pass_context def train_cmd(ctx, path, input, output, keep_venv): """Execute 'train' entrypoint in fedn.yaml. @@ -108,11 +107,11 @@ def train_cmd(ctx, path, input, output, keep_venv): dispatcher = Dispatcher(config, path) _ = dispatcher._get_or_create_python_env() dispatcher.run_cmd("train {} {}".format(input, output)) - delete_virtual_environment(keep_venv, dispatcher) - + if not keep_venv: + delete_virtual_environment(dispatcher) @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.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set doesn't remove venv") @click.pass_context def startup_cmd(ctx, path, keep_venv): """Execute 'startup' entrypoint in fedn.yaml. @@ -132,12 +131,12 @@ def startup_cmd(ctx, path, keep_venv): dispatcher = Dispatcher(config, path) _ = dispatcher._get_or_create_python_env() dispatcher.run_cmd("startup") - delete_virtual_environment(keep_venv, dispatcher) - + if not keep_venv: + delete_virtual_environment(dispatcher) @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.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set doesn't remove venv") @click.pass_context def build_cmd(ctx, path, keep_venv): """Execute 'build' entrypoint in fedn.yaml. @@ -158,8 +157,8 @@ def build_cmd(ctx, path, keep_venv): dispatcher = Dispatcher(config, path) _ = dispatcher._get_or_create_python_env() dispatcher.run_cmd("build") - delete_virtual_environment(keep_venv, dispatcher) - + if not keep_venv: + delete_virtual_environment(dispatcher) @run_cmd.command("client")