Refactor Code For Better Extensibility

deployment/docker_trigger.py

@@ -4,10 +4,10 @@
 from echoflow import echoflow_start
 from echoflow.stages.echoflow_trigger import echoflow_trigger
 from prefect import flow
-from prefect.task_runners import SequentialTaskRunner
+from prefect.task_runners import ThreadPoolTaskRunner
 from typing import Any, Dict, Optional, Union
 
-@flow(name="Docker-Trigger", task_runner=SequentialTaskRunner())
+@flow(name="Docker-Trigger", task_runner=ThreadPoolTaskRunner(max_workers=1))
 def docker_trigger(
     dataset_config: Union[dict, str, Path],
     pipeline_config: Union[dict, str, Path],

echodataflow/__init__.py

@@ -7,8 +7,9 @@
                                        echodataflow_create_prefect_profile,
                                        echodataflow_start, get_active_profile,
                                        load_profile)
-from .utils.config_utils import extract_fs, glob_url, load_block
+from .utils.config_utils import extract_fs, glob_url
 from .utils.file_utils import get_ed_list, get_last_run_output, get_zarr_list
+from .utils.filesystem_utils import load_block
 from .docker_trigger import docker_trigger
 
 try:

echodataflow/docker_trigger.py

@@ -1,10 +1,10 @@
 from pathlib import Path
 from echodataflow.stages.echodataflow_trigger import echodataflow_trigger
 from prefect import flow
-from prefect.task_runners import SequentialTaskRunner
+from prefect.task_runners import ThreadPoolTaskRunner
 from typing import Any, Dict, Optional, Union
 
-@flow(name="docker-trigger-latest", task_runner=SequentialTaskRunner())
+@flow(name="docker-trigger-latest", task_runner=ThreadPoolTaskRunner(max_workers=1))
 def docker_trigger(
     dataset_config: Union[dict, str, Path],
     pipeline_config: Union[dict, str, Path],

echodataflow/extensions/file_downloader.py

@@ -15,10 +15,12 @@
 from prefect.client.schemas.objects import FlowRun, StateType
 from prefect.states import Cancelled
 
-from echodataflow.utils.config_utils import glob_url, handle_storage_options
+from echodataflow.utils.config_utils import glob_url
 from echodataflow.utils.file_utils import extract_fs, make_temp_folder
 import shlex
 
+from echodataflow.utils.filesystem_utils import handle_storage_options
+
 
 @task
 def download_temp_file(file_url: str, storage_options: Dict[str, Any], dest_dir: str, delete_on_transfer: bool, replace: bool) -> str:

echodataflow/extensions/file_monitor.py

@@ -19,8 +19,8 @@
 
 from echodataflow.models.datastore import StorageType
 from echodataflow.models.run import EDFRun, FileDetails
-from echodataflow.utils.config_utils import glob_url, load_block
-from prefect.task_runners import SequentialTaskRunner
+from echodataflow.utils.config_utils import glob_url
+from echodataflow.utils.filesystem_utils import load_block
 
 @task
 def execute_flow(
@@ -112,10 +112,10 @@ def file_monitor(
 
     new_run = datetime.now(tz=timezone.utc).isoformat()
     edfrun: EDFRun = None
-    try:
+    try:        
         edfrun = load_block(
                     name=block_name,
-                    type=StorageType.EDFRUN,
+                    stype=StorageType.EDFRUN,
                 )
     except Exception as e:
         print(e)        

echodataflow/models/echodataflow_config.py

@@ -18,11 +18,9 @@
 Email: [email protected]
 Date: August 22, 2023
 """
-import json
 from typing import Any, Dict, List, Optional
 
 from prefect.blocks.core import Block
-from pydantic import SecretStr
 
 from .datastore import StorageType
 
@@ -43,9 +41,9 @@ class EchodataflowPrefectConfig(Block):
     class Config:
         arbitrary_types_allowed = True
 
-    prefect_account_id: str = None
-    prefect_api_key: str = None
-    prefect_workspace_id: str = None
+    prefect_account_id: Optional[str] = None
+    prefect_api_key: Optional[str] = None
+    prefect_workspace_id: Optional[str] = None
     profile_name: str = None
 
 

echodataflow/stages/echodataflow.py

@@ -47,9 +47,9 @@
 )
 
 import echopype as ep
-from echodataflow.utils.config_utils import load_block
 
 from echodataflow.stages.echodataflow_trigger import echodataflow_trigger
+from echodataflow.utils.filesystem_utils import handle_storage_options
 
 
 def check_internet_connection(host="8.8.8.8", port=53, timeout=5):
@@ -236,7 +236,7 @@ def echodataflow_start(
 
     # Try loading the Prefect config block
     try:
-        load_block(name="echodataflow-config", type=StorageType.ECHODATAFLOW)
+        handle_storage_options({'block_name':"echodataflow-config", 'type':StorageType.ECHODATAFLOW})        
     except ValueError:
         print(
             "\nNo Prefect Cloud Configuration found. Creating Prefect Local named 'echodataflow-local'. Please add your prefect cloud "
@@ -546,8 +546,7 @@ def load_credential_configuration(sync: bool = False):
                 current_config = asyncio.run(current_config)
             if current_config is not None:
                 for base in current_config.blocks:
-                    block = load_block(base.name, base.type)
-                    block_dict = dict(block)
+                    block_dict = handle_storage_options(base)  
                     block_dict["name"] = base.name
                     block_dict["active"] = base.active
                     block_dict["options"] = json.dumps(base.options)

echodataflow/stages/echodataflow_trigger.py

@@ -17,28 +17,24 @@
 import json
 from pathlib import Path
 from typing import Optional, Union
-from fastapi.encoders import jsonable_encoder
 
+from fastapi.encoders import jsonable_encoder
 from prefect import flow
-from prefect.task_runners import SequentialTaskRunner
-from prefect.blocks.core import Block
-from prefect.variables import Variable
+from prefect.task_runners import ThreadPoolTaskRunner
 
 from echodataflow.aspects.singleton_echodataflow import Singleton_Echodataflow
 from echodataflow.models.datastore import Dataset
 from echodataflow.models.pipeline import Recipe
 from echodataflow.utils import log_util
-from echodataflow.utils.config_utils import (
-    check_config,
-    extract_config,
-    get_storage_options,
-    load_block,
-)
+from echodataflow.utils.config_utils import (check_config,
+                                             parse_dynamic_parameters,
+                                             parse_yaml_config)
+from echodataflow.utils.filesystem_utils import handle_storage_options
 
 from .subflows.initialization_flow import init_flow
 
 
-@flow(name="Echodataflow", task_runner=SequentialTaskRunner())
+@flow(name="Echodataflow", task_runner=ThreadPoolTaskRunner(max_workers=1))
 def echodataflow_trigger(
     dataset_config: Union[dict, str, Path],
     pipeline_config: Union[dict, str, Path],
@@ -82,48 +78,15 @@ def echodataflow_trigger(
         print("Pipeline output:", pipeline_output)
     """
 
-    if storage_options:
-        # Check if storage_options is a Block (fsspec storage) and convert it to a dictionary
-        if isinstance(storage_options, Block):
-            storage_options = get_storage_options(storage_options=storage_options)
-        elif isinstance(storage_options, dict) and storage_options.get("block_name"):
-            block = load_block(
-                name=storage_options.get("block_name"), type=storage_options.get("type")
-            )
-            storage_options = get_storage_options(block)
-    else:
-        storage_options = {}
-
-    if isinstance(dataset_config, Path):
-        dataset_config = str(dataset_config)
-    if isinstance(logging_config, Path):
-        logging_config = str(logging_config)
-    if isinstance(pipeline_config, Path):
-        pipeline_config = str(pipeline_config)
+    storage_options = handle_storage_options(storage_options=storage_options)
+
+    dataset_config_dict = parse_yaml_config(config=dataset_config, storage_options=storage_options)
+    logging_config_dict = parse_yaml_config(config=logging_config, storage_options=storage_options)
+    pipeline_config_dict = parse_yaml_config(config=pipeline_config, storage_options=storage_options)
+
     if isinstance(json_data_path, Path):
         json_data_path = str(json_data_path)
 
-    if isinstance(dataset_config, str):
-        if not dataset_config.endswith((".yaml", ".yml")):
-            raise ValueError("Configuration file must be a YAML!")
-        dataset_config_dict = extract_config(dataset_config, storage_options)
-    elif isinstance(dataset_config, dict):
-        dataset_config_dict = dataset_config
-
-    if isinstance(pipeline_config, str):
-        if not pipeline_config.endswith((".yaml", ".yml")):
-            raise ValueError("Configuration file must be a YAML!")
-        pipeline_config_dict = extract_config(pipeline_config, storage_options)
-    elif isinstance(pipeline_config, dict):
-        pipeline_config_dict = pipeline_config
-
-    if isinstance(logging_config, str):
-        if not logging_config.endswith((".yaml", ".yml")):
-            raise ValueError("Configuration file must be a YAML!")
-        logging_config_dict = extract_config(logging_config, storage_options)
-    else:
-        logging_config_dict = logging_config
-
     log_util.log(
         msg={
             "msg": f"Dataset Configuration Loaded For This Run",
@@ -144,7 +107,6 @@ def echodataflow_trigger(
         },
         eflogging=dataset_config_dict.get("logging"),
     )
-    print(dataset_config_dict)
 
     log_util.log(
         msg={
@@ -172,62 +134,18 @@ def echodataflow_trigger(
     check_config(dataset_config_dict, pipeline_config_dict)
     pipeline = Recipe(**pipeline_config_dict)
     dataset = Dataset(**dataset_config_dict)
-
-    if options.get("storage_options_override") and not options["storage_options_override"]:
-        storage_options = {}
-    if not storage_options:
-        if dataset.output.storage_options:
-            if not dataset.output.storage_options.anon:
-                block = load_block(
-                    name=dataset.output.storage_options.block_name,
-                    type=dataset.output.storage_options.type,
-                )
-                dataset.output.storage_options_dict = get_storage_options(block)
-            else:
-                dataset.output.storage_options_dict = {"anon": dataset.output.storage_options.anon}
-
-        if dataset.args.storage_options:
-            if not dataset.args.storage_options.anon:
-                block = load_block(
-                    name=dataset.args.storage_options.block_name,
-                    type=dataset.args.storage_options.type,
-                )
-                dataset.args.storage_options_dict = get_storage_options(block)
-            else:
-                dataset.args.storage_options_dict = {"anon": dataset.args.storage_options.anon}
-        if dataset.args.group:
-            if dataset.args.group.storage_options:
-                if not dataset.args.group.storage_options.anon:
-                    block = load_block(
-                        name=dataset.args.group.storage_options.block_name,
-                        type=dataset.args.group.storage_options.type,
-                    )
-                    dataset.args.group.storage_options_dict = get_storage_options(block)
-                else:
-                    dataset.args.group.storage_options_dict = {
-                        "anon": dataset.args.group.storage_options.anon
-                    }
-    else:
+
+    if options.get("storage_options_override", False):
         dataset.output.storage_options_dict = storage_options
         dataset.args.storage_options_dict = storage_options
         dataset.args.group.storage_options_dict = storage_options
+    else:
+        dataset.output.storage_options_dict = handle_storage_options(storage_options=dataset.output.storage_options)
+        dataset.args.storage_options_dict = handle_storage_options(storage_options=dataset.args.storage_options)
+        dataset.args.group.storage_options_dict = handle_storage_options(storage_options=dataset.args.group.storage_options)
+
+    edf = Singleton_Echodataflow(log_file=logging_config_dict, pipeline=pipeline, dataset=dataset)
 
-    print("\nInitiliazing Singleton Object")
-    Singleton_Echodataflow(log_file=logging_config_dict, pipeline=pipeline, dataset=dataset)
-
-    if dataset.args.parameters and dataset.args.parameters.file_name and dataset.args.parameters.file_name == "VAR_RUN_NAME":
-        var: Variable = Variable.get("run_name", default=None)
-        if not var:
-            raise ValueError("No variable found for name `run_name`")
-        else:
-            dataset.args.parameters.file_name = var.value
-
-    # Change made to enable dynamic execution using an extension
-    if options and options.get("file_name"):
-        dataset.args.parameters.file_name = options.get("file_name")
+    dataset = parse_dynamic_parameters(dataset, options=options)
 
-    if options and options.get("run_name"):
-        dataset.name = options.get("run_name")
-
-    print("\nReading Configurations")
     return init_flow(config=dataset, pipeline=pipeline, json_data_path=json_data_path)