<gh_stars>0
#!/usr/bin/env python

####################
# Required Modules #
####################

# Generic/Built-in
import asyncio
import json
import os
import shutil
from pathlib import Path
from typing import Dict, List, Any

# Libs
import aiohttp
import mlflow
import torch as th

# Custom
from rest_rpc import app
from rest_rpc.training.core.utils import (
    UrlConstructor, 
    RPCFormatter, 
    Orchestrator
)
from synarchive.connection import ExperimentRecords, RunRecords
from synarchive.training import ModelRecords
from synarchive.evaluation import MLFRecords

##################
# Configurations #
##################

SOURCE_FILE = os.path.abspath(__file__)

schemas = app.config['SCHEMAS']
db_path = app.config['DB_PATH']
out_dir = app.config['OUT_DIR']
mlflow_dir = app.config['MLFLOW_DIR']

logging = app.config['NODE_LOGGER'].synlog
logging.debug("evaluation/core/utils.py logged", Description="No Changes")

####################
# Helper Functions #
####################

def replicate_combination_key(expt_id, run_id):
    return str((expt_id, run_id))

############################################
# Inference Orchestration class - Analyser #
############################################

class Analyser(Orchestrator):
    """ 
    Takes in a list of minibatch IDs and sends them to worker nodes. Workers
    will use these IDs to reconstruct their aggregated test datasets with 
    prediction labels mapped appropriately.

    Attributes:
        inferences (dict(str, dict(str, dict(str, th.tensor)))
    """
    def __init__(
        self, 
        collab_id: str,
        project_id: str,
        expt_id: str, 
        run_id: str,
        inferences: dict,
        metas: list = ['train', 'evaluate', 'predict'],
        auto_align: bool = True
    ):
        super().__init__()

        self.metas = metas
        self.collab_id = collab_id
        self.project_id = project_id
        self.expt_id = expt_id
        self.run_id = run_id
        self.inferences = inferences
        self.auto_align = auto_align

    ###########
    # Helpers #
    ###########

    async def _poll_for_stats(
        self, 
        node_info: Dict[str, Any], 
        inferences: Dict[str, Dict[str, th.Tensor]]
    ):
        """ Parses a registration record for participant metadata, before
            submitting minibatch IDs of inference objects to corresponding 
            worker node's REST-RPC service for calculating descriptive
            statistics and prediction exports

        Args:
            reg_record (tinydb.database.Document): Participant-project details
            inferences (dict): List of dicts containing inference object IDs
        Returns:
            Statistics (dict)
        """
        _, _, participant_id = self.parse_keys(node_info)
 
        if not inferences:
            return {participant_id: {meta:{} for meta in self.metas}}

        # Construct destination url for interfacing with worker REST-RPC
        rest_connection = self.parse_rest_info(node_info)
        destination_constructor = UrlConstructor(**rest_connection)
        destination_url = destination_constructor.construct_predict_url(
            collab_id=self.collab_id,
            project_id=self.project_id,
            expt_id=self.expt_id,
            run_id=self.run_id
        )

        ml_action = self.parse_action(node_info)
        data_tags = self.parse_tags(node_info)
        data_alignments = self.parse_alignments(node_info, self.auto_align)

        payload = {
            'action': ml_action, 
            'tags': data_tags,
            'alignments': data_alignments,
            'inferences': inferences
        }

        # Trigger remote inference by posting alignments & ID mappings to 
        # `Predict` route in worker
        resp_inference_data, status_code = await self.instruct(
            command='post', 
            url=destination_url, 
            payload=payload
        )
        
        logging.debug(
            f"Participant '{participant_id}' >|< Project '{self.project_id}' -> Experiment '{self.expt_id}' -> Run '{self.run_id}': Polled statistics tracked.",
            description=f"Polled statistics for participant '{participant_id}' under project '{self.project_id}' using experiment '{self.expt_id}' and run '{self.run_id}' tracked.",
            resp_json=resp_inference_data,
            ID_path=SOURCE_FILE,
            ID_class=Analyser.__name__,
            ID_function=Analyser._poll_for_stats.__name__
        )

        # Extract the relevant expt-run results
        expt_run_key = replicate_combination_key(self.expt_id, self.run_id)
        metadata = resp_inference_data['results'][expt_run_key]

        # Filter by selected meta-datasets
        filtered_statistics = {
            meta: stats 
            for meta, stats in metadata.items()
            if meta in self.metas
        }

        return {participant_id: filtered_statistics}


    async def _collect_all_stats(self, grid: List[Dict[str, Any]]) -> dict:
        """ Asynchronous function to submit inference data to registered
            participant servers in return for remote performance statistics

        Args:
            grid (list(dict))): Registry of participants' node information
        Returns:
            All participants' statistics (dict)
        """

        ###########################
        # Implementation Footnote #
        ###########################

        # [Causes]
        # Node info retrieved is asynmetric as compared to input inferences.

        # [Problems]
        # Zipping causes misallocation of dataset inferences to the wrong
        # servers, resulting in a "ValueError: Found input variables with 
        # inconsistent numbers of samples: [N1, N2]", since prediction counts 
        # obtained are not aligned with label counts in the remote servers 

        # [Solution]
        # Pair things up via manual matching

        mapped_pairs = [
            (record, inferences)
            for worker_id, inferences in self.inferences.items()
            for record in grid
            if self.parse_syft_info(record).get('id') == worker_id
        ]

        logging.debug(f"Mapped pairs: {mapped_pairs}")

        all_statistics = {}
        for future in asyncio.as_completed(
            map(lambda args: self._poll_for_stats(*args), mapped_pairs)
        ):
            result = await future
            all_statistics.update(result)

        return all_statistics

    ##################
    # Core Functions #
    ##################

    def infer(self, grid: List[Dict[str, Any]]) -> dict:
        """ Wrapper function for triggering asychroneous remote inferencing of
            participant nodes

        Args:
            grid (list(dict))): Registry of participants' node information
        Returns:
            All participants' statistics (dict)
        """
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        try:
            all_stats = loop.run_until_complete(
                self._collect_all_stats(grid=grid)
            )
        finally:
            loop.close()

        return all_stats



####################################
# MLFLow logging class - MLFlogger #
####################################

class MLFlogger:
    """ 
    Wrapper class around MLFlow to faciliate experiment & run registrations in
    the REST-RPC setting, where statistical logging is performed during 
    post-mortem analysis
    """
    def __init__(self, db_path: str = db_path, remote_uri: str = None):

        # Private attributes
        self.__rpc_formatter = RPCFormatter()
        self.__expt_records = ExperimentRecords(db_path=db_path)
        self.__run_records = RunRecords(db_path=db_path)
        self.__model_records = ModelRecords(db_path=db_path)
        self.__connector = "_>_"
        
        # Public attributes
        self.mlf_records = MLFRecords(db_path=db_path)
        self.db_path = db_path
        self.remote_uri = remote_uri

    ###########
    # Helpers #
    ###########

    def _generate_expt_name(
        self, 
        collab_id: str, 
        project_id: str, 
        expt_id: str
    ) -> str:
        """ Abstracts creation of a unique key to identify an MLFlow experiment. 
            This collapses the payload hierarchy so that only one MLFlow instance
            is required to host all runs registered under collaboration(s)  

        Args:
            collab_id (str): ID of collaboration
            project_id (str): ID of project
            expt_id (str): ID of experiment
        Returns:
            Experiment name (str)
        """
        return self.__connector.join([collab_id, project_id, expt_id])


    def _generate_run_name(
        self,         
        collab_id: str, 
        project_id: str, 
        expt_id: str,
        run_id: str
    ) -> str:
        """ Abstracts creation of a unique key to identify an MLFlow run.

        Args:
            collab_id (str): ID of collaboration
            project_id (str): ID of project
            expt_id (str): ID of experiment
            run_id (str): ID of run
        Returns:
            Run name (str)
        """
        return run_id   # keeping it simple

    
    def _generate_record_name(
        self,
        collab_id: str, 
        project_id: str, 
        expt_id: str,
        run_id: str = "",
        mlflow_type: str = "experiment"
    ):
        """ Abstracts creation of a unique key to partition experiments & runs
            internally within the context of MLFlow.

        Args:
            collab_id (str): ID of collaboration
            project_id (str): ID of project
            expt_id (str): ID of experiment
            run_id (str): ID of run
            mlflow_type (str): Type of entity to be represented in MLFlow
        Returns:
            Record name (str)
        """
        
        ###########################
        # Implementation Footnote #
        ###########################

        # [Causes]
        # In MLFlow, only experiments and runs are mappable/usable to the 
        # Synergos hierarchy. In order to simplify the deployment process, 
        # collaborations & projects are collapsed in as part of an MLFlow's 
        # experiment, allowing a single instance of MLFlow to host results &
        # analysis of results across the grid.

        # [Problems]
        # The collapsed hierarchy becomes ambigious due to the inability to
        # distinguish between a Synergos experiment and a Synergos run.

        # [Solution]
        # Use expt_id and run_id to create a unique composite key for use when
        # storing mapping details between MLFlow and Synergos.

        if mlflow_type == "experiment":
            return expt_id

        elif mlflow_type == "run":
            return self.__connector.join([expt_id, run_id])

        else:
            raise RuntimeError("Unsupported entity specified!")


    def retrieve_mlflow_experiment(
        self,
        collab_id: str, 
        project_id: str,
        expt_id: str,
    ) -> Dict[str, str]:
        """
        """
        expt_name = self._generate_expt_name(collab_id, project_id, expt_id)
        expt_record_name = self._generate_record_name(
            collab_id, project_id, expt_id,
            mlflow_type="experiment"
        )
        expt_mlflow_record = self.mlf_records.read(
            collaboration=collab_id,
            project=project_id, 
            record=expt_record_name,
            name=expt_name
        )
        return expt_mlflow_record


    def retrieve_mlflow_run(
        self,
        collab_id: str, 
        project_id: str,
        expt_id: str,
        run_id: str
    ) -> Dict[str, str]:
        """
        """
        run_name = self._generate_run_name(collab_id, project_id, expt_id, run_id)
        run_record_name = self._generate_record_name(
            collab_id, project_id, expt_id, run_id,
            mlflow_type="run"
        )
        run_mlflow_record = self.mlf_records.read(
            collaboration=collab_id,
            project=project_id, 
            record=run_record_name,
            name=run_name
        )
        return run_mlflow_record


    def initialise_mlflow_experiment(
        self, 
        collab_id: str, 
        project_id: str,
        expt_id: str,
        tracking_uri: str
    ) -> Dict[str, str]:
        """ Initialises an MLFlow experiment at the specified project URI

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
        Returns:
            MLFlow Experiment configuration (dict)
        """
        mlflow.set_tracking_uri(tracking_uri)

        expt_name = self._generate_expt_name(collab_id, project_id, expt_id)
        expt_record_name = self._generate_record_name(
            collab_id, project_id, expt_id,
            mlflow_type="experiment"
        )
        expt_mlflow_record = self.mlf_records.read(
            collaboration=collab_id,
            project=project_id, 
            record=expt_record_name,
            name=expt_name
        )

        # Check if MLFlow experiment has already been created. Create a 
        # MLFRecord if: 
        # 1) No previous records for this combination key was initialized
        # 2) Specified tracking URI has changed

        if not expt_mlflow_record:
            
            # Initialise MLFlow experiment
            mlflow_id = mlflow.create_experiment(name=expt_name)

            mlflow_details = {
                'collaboration': collab_id,
                'project': project_id,
                'record': expt_record_name,
                'name': expt_name,
                'mlflow_type': 'experiment',
                'mlflow_id': mlflow_id,
                'mlflow_uri': tracking_uri  # Overridable with remote > local
            }
            expt_mlflow_record = self.mlf_records.create(
                collaboration=collab_id,
                project=project_id,
                record=expt_record_name,
                name=expt_name, 
                details=mlflow_details
            )

        # Update URI if necessary
        if expt_mlflow_record.get('mlflow_uri') != tracking_uri:
            expt_mlflow_record = self.mlf_records.update(
                collaboration=collab_id,
                project=project_id,
                record=expt_record_name,
                name=expt_name, 
                updates={'mlflow_uri': tracking_uri}
            )

        stripped_expt_mlflow_details = self.__rpc_formatter.strip_keys(
            record=expt_mlflow_record
        )
        return stripped_expt_mlflow_details


    def initialise_mlflow_run(
        self, 
        collab_id: str, 
        project_id: str, 
        expt_id: str, 
        run_id: str
    ) -> Dict[str, str]:
        """ Initialises a MLFLow run under a specified experiment of a project.
            Initial run hyperparameters will be logged, and MLFlow run id will
            be stored for subsequent analysis.

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
            run_id (str): REST-RPC ID of specified run
        """
        # Retrieve parent MLFlow experiment metadata (if available)
        expt_mlflow_record = self.retrieve_mlflow_experiment(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id
        )
        expt_mlflow_id = expt_mlflow_record.get('mlflow_id')
        expt_mlflow_uri = expt_mlflow_record.get('mlflow_uri')

        if not expt_mlflow_id:
            logging.error(
                "MLFlow experiment has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.initialise_mlflow_run.__name__
            )
            raise RuntimeError("MLFlow experiment has not been initialised!")

        run_name = self._generate_run_name(collab_id, project_id, expt_id, run_id)
        run_record_name = self._generate_record_name(
            collab_id, project_id, expt_id, run_id,
            mlflow_type="run"
        )
        run_mlflow_record = self.mlf_records.read(
            collaboration=collab_id,
            project=project_id, 
            record=run_record_name,
            name=run_name
        )
        if not run_mlflow_record:

            with mlflow.start_run(
                experiment_id=expt_mlflow_id, 
                run_name=run_name
            ) as mlf_run:

                # Retrieve run details from database
                run_details = self.__run_records.read(
                    collab_id=collab_id, 
                    project_id=project_id, 
                    expt_id=expt_id, 
                    run_id=run_id
                )
                stripped_run_details = self.__rpc_formatter.strip_keys(
                    record=run_details,
                    concise=True
                )

                mlflow.log_params(stripped_run_details)

                # Save the MLFlow ID mapping
                run_mlflow_id = mlf_run.info.run_id
                run_mlflow_details = {
                    'collaboration': collab_id,
                    'project': project_id,
                    'record': run_record_name,
                    'name': run_name,
                    'mlflow_type': 'run',
                    'mlflow_id': run_mlflow_id,
                    'mlflow_uri': expt_mlflow_record['mlflow_uri'] # same as expt
                }
                run_mlflow_record = self.mlf_records.create(
                    collaboration=collab_id,
                    project=project_id,
                    record=run_record_name,
                    name=run_name, 
                    details=run_mlflow_details
                )

        if run_mlflow_record.get('mlflow_uri') != expt_mlflow_uri:
            run_mlflow_record = self.mlf_records.update(
                collaboration=collab_id,
                project=project_id,
                record=run_record_name,
                name=run_name, 
                updates={'mlflow_uri': expt_mlflow_record['mlflow_uri']}
            )

        stripped_run_mlflow_details = self.__rpc_formatter.strip_keys(
            record=run_mlflow_record
        )
        return stripped_run_mlflow_details       


    def log_losses(
        self, 
        collab_id: str, 
        project_id: str, 
        expt_id: str, 
        run_id: str
    ):
        """ Registers all cached losses, be it global or local, obtained from
            federated training into MLFlow.

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
            run_id (str): REST-RPC ID of specified run
        Returns:
            Stripped metadata (dict)
        """
        # Retrieve parent MLFlow experiment metadata (if available)
        expt_mlflow_record = self.retrieve_mlflow_experiment(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id
        )
        expt_mlflow_id = expt_mlflow_record.get('mlflow_id')

        if not expt_mlflow_id:
            logging.error(
                "MLFlow experiment has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow experiment has not been initialised!")

        # Search for run session to update entry, not create a new one
        run_mlflow_record = self.retrieve_mlflow_run(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id, 
            run_id=run_id
        )
        run_mlflow_id = run_mlflow_record.get('mlflow_id')

        if not run_mlflow_id:
            logging.error(
                "MLFlow run has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow run has not been initialised!")

        # Retrieve all model metadata from storage
        model_metadata = self.__model_records.read(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id, 
            run_id=run_id
        )
        stripped_metadata = self.__rpc_formatter.strip_keys(
            record=model_metadata, 
            concise=True
        ) if model_metadata else model_metadata

        if stripped_metadata:

            with mlflow.start_run(
                experiment_id=expt_mlflow_id, 
                run_id=run_mlflow_id # continue existing run
            ) as mlf_run:

                # Extract loss histories
                for m_type, metadata in stripped_metadata.items():

                    loss_history_path = metadata['loss_history']
                    origin = metadata['origin']

                    with open(loss_history_path, 'r') as lh:
                        loss_history = json.load(lh)

                    if m_type == 'global':
                        for meta, losses in loss_history.items():
                            for round_idx, loss, in losses.items():
                                mlflow.log_metric(
                                    key=f"global_{meta}_loss", 
                                    value=loss, 
                                    step=int(round_idx)
                                )

                    else:
                        for round_idx, loss, in losses.items():
                            mlflow.log_metric(
                                key=f"{origin}_local_loss", 
                                value=loss, 
                                step=int(round_idx)
                            )

        return stripped_metadata

    
    def log_model_performance(
        self, 
        collab_id: str,
        project_id: str, 
        expt_id: str, 
        run_id: str,
        statistics: dict
    ) -> dict:
        """ Using all cached model checkpoints, log the performance statistics 
            of models, be it global or local, to MLFlow at round (for global) or
            epoch (for local) level.

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
            run_id (str): REST-RPC ID of specified run
            statistics (dict): Inference statistics polled from workers
        Returns:
            MLFLow run details (dict)
        """
        # Retrieve parent MLFlow experiment metadata (if available)
        expt_mlflow_record = self.retrieve_mlflow_experiment(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id
        )
        expt_mlflow_id = expt_mlflow_record.get('mlflow_id')

        if not expt_mlflow_id:
            logging.error(
                "MLFlow experiment has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow experiment has not been initialised!")

        # Search for run session to update entry, not create a new one
        run_mlflow_record = self.retrieve_mlflow_run(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id, 
            run_id=run_id
        )
        run_mlflow_id = run_mlflow_record.get('mlflow_id')

        if not run_mlflow_id:
            logging.error(
                "MLFlow run has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow run has not been initialised!")

        with mlflow.start_run(
            experiment_id=expt_mlflow_id, 
            run_id=run_mlflow_id
        ) as mlf_run:

            # Store output metadata into database
            for _, inference_stats in statistics.items():
                for _, meta_stats in inference_stats.items():

                    # Log statistics to MLFlow for analysis
                    stats = meta_stats.get('statistics', {})
                    for stat_type, stat_value in stats.items():

                        if isinstance(stat_value, list):
                            for val_idx, value in enumerate(stat_value):
                                mlflow.log_metric(
                                    key=f"{stat_type}_class_{val_idx}", 
                                    value=value, 
                                    step=int(val_idx+1)
                                )

                        else:
                            mlflow.log_metric(key=stat_type, value=stat_value)

        stripped_mlflow_run_details = self.__rpc_formatter.strip_keys(
            run_mlflow_record, 
            concise=True
        )
        return stripped_mlflow_run_details


    def log_artifacts(
        self, 
        collab_id: str,
        project_id: str, 
        expt_id: str, 
        run_id: str
    ) -> str:
        """ Log all artifacts produced (i.e. local & global models, losses etc.)
            during the specified federated run.

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
            run_id (str): REST-RPC ID of specified run
            statistics (dict): Inference statistics polled from workers
        Returns:
            MLFLow run details (dict)
        """
        # Retrieve parent MLFlow experiment metadata (if available)
        expt_mlflow_record = self.retrieve_mlflow_experiment(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id
        )
        expt_mlflow_id = expt_mlflow_record.get('mlflow_id')

        if not expt_mlflow_id:
            logging.error(
                "MLFlow experiment has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow experiment has not been initialised!")

        # Search for run session to update entry, not create a new one
        run_mlflow_record = self.retrieve_mlflow_run(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id, 
            run_id=run_id
        )
        run_mlflow_id = run_mlflow_record.get('mlflow_id')

        if not run_mlflow_id:
            logging.error(
                "MLFlow run has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow run has not been initialised!")
       
        with mlflow.start_run(
            experiment_id=expt_mlflow_id, 
            run_id=run_mlflow_id # continue existing run
        ) as mlf_run:

            results_dir = os.path.join(out_dir, collab_id, project_id, expt_id, run_id)           
            mlflow.log_artifacts(local_dir=results_dir)

        return results_dir

    ##################
    # Core Functions #
    ##################

    def log(self, accumulations: dict) -> List[str]:
        """ Wrapper function that processes statistics accumulated from 
            inference.

        Args:
            accumulations (dict): Accumulated statistics from inferring
                different project-expt-run combinations
        Returns:
            List of MLFlow run IDs from all runs executed (list(str))
        """

        ###########################
        # Implementation Footnote #
        ###########################

        # [Causes]
        # Depending on where your MLFlow logs are, the accessibility of the
        # resource changes.

        # [Problems]
        # As Synergos MLOps is an optional component, the orchestrator may or
        # may not install it. Also, depending on the location of deployment
        # (i.e. local filesystem vs remote tracking server), connection URI
        # may vary.

        # [Solution]
        # By default, ALWAYS generate local MLFlow caches. If a remote tracking
        # URI was specified, then attempt to send logs over. That way, clients
        # will always have the option of choosing which way they would to serve
        # their MLOps information

        has_remote_tracking = self.remote_uri is not None
        tracking_uris = (
            [mlflow_dir, self.remote_uri] 
            if has_remote_tracking 
            else [mlflow_dir]
        ) 

        jobs_ran = []
        for tracking_uri in tracking_uris:
            for combination_key, statistics in accumulations.items():

                collab_id, project_id, expt_id, run_id = combination_key

                self.initialise_mlflow_experiment(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    tracking_uri=tracking_uri
                )
                run_mlflow_details = self.initialise_mlflow_run(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    run_id=run_id
                )
                self.log_losses(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    run_id=run_id
                )
                self.log_model_performance(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    run_id=run_id,
                    statistics=statistics
                )
                self.log_artifacts(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    run_id=run_id
                )

                if run_mlflow_details['mlflow_id'] not in jobs_ran:
                    jobs_ran.append(run_mlflow_details['mlflow_id'])
        
        return jobs_ran
#!/usr/bin/env python

####################
# Required Modules #
####################

# Generic/Built-in
import asyncio
import json
import os
import shutil
from pathlib import Path
from typing import Dict, List, Any

# Libs
import aiohttp
import mlflow
import torch as th

# Custom
from rest_rpc import app
from rest_rpc.training.core.utils import (
    UrlConstructor, 
    RPCFormatter, 
    Orchestrator
)
from synarchive.connection import ExperimentRecords, RunRecords
from synarchive.training import ModelRecords
from synarchive.evaluation import MLFRecords

##################
# Configurations #
##################

SOURCE_FILE = os.path.abspath(__file__)

schemas = app.config['SCHEMAS']
db_path = app.config['DB_PATH']
out_dir = app.config['OUT_DIR']
mlflow_dir = app.config['MLFLOW_DIR']

logging = app.config['NODE_LOGGER'].synlog
logging.debug("evaluation/core/utils.py logged", Description="No Changes")

####################
# Helper Functions #
####################

def replicate_combination_key(expt_id, run_id):
    return str((expt_id, run_id))

############################################
# Inference Orchestration class - Analyser #
############################################

class Analyser(Orchestrator):
    """ 
    Takes in a list of minibatch IDs and sends them to worker nodes. Workers
    will use these IDs to reconstruct their aggregated test datasets with 
    prediction labels mapped appropriately.

    Attributes:
        inferences (dict(str, dict(str, dict(str, th.tensor)))
    """
    def __init__(
        self, 
        collab_id: str,
        project_id: str,
        expt_id: str, 
        run_id: str,
        inferences: dict,
        metas: list = ['train', 'evaluate', 'predict'],
        auto_align: bool = True
    ):
        super().__init__()

        self.metas = metas
        self.collab_id = collab_id
        self.project_id = project_id
        self.expt_id = expt_id
        self.run_id = run_id
        self.inferences = inferences
        self.auto_align = auto_align

    ###########
    # Helpers #
    ###########

    async def _poll_for_stats(
        self, 
        node_info: Dict[str, Any], 
        inferences: Dict[str, Dict[str, th.Tensor]]
    ):
        """ Parses a registration record for participant metadata, before
            submitting minibatch IDs of inference objects to corresponding 
            worker node's REST-RPC service for calculating descriptive
            statistics and prediction exports

        Args:
            reg_record (tinydb.database.Document): Participant-project details
            inferences (dict): List of dicts containing inference object IDs
        Returns:
            Statistics (dict)
        """
        _, _, participant_id = self.parse_keys(node_info)
 
        if not inferences:
            return {participant_id: {meta:{} for meta in self.metas}}

        # Construct destination url for interfacing with worker REST-RPC
        rest_connection = self.parse_rest_info(node_info)
        destination_constructor = UrlConstructor(**rest_connection)
        destination_url = destination_constructor.construct_predict_url(
            collab_id=self.collab_id,
            project_id=self.project_id,
            expt_id=self.expt_id,
            run_id=self.run_id
        )

        ml_action = self.parse_action(node_info)
        data_tags = self.parse_tags(node_info)
        data_alignments = self.parse_alignments(node_info, self.auto_align)

        payload = {
            'action': ml_action, 
            'tags': data_tags,
            'alignments': data_alignments,
            'inferences': inferences
        }

        # Trigger remote inference by posting alignments & ID mappings to 
        # `Predict` route in worker
        resp_inference_data, status_code = await self.instruct(
            command='post', 
            url=destination_url, 
            payload=payload
        )
        
        logging.debug(
            f"Participant '{participant_id}' >|< Project '{self.project_id}' -> Experiment '{self.expt_id}' -> Run '{self.run_id}': Polled statistics tracked.",
            description=f"Polled statistics for participant '{participant_id}' under project '{self.project_id}' using experiment '{self.expt_id}' and run '{self.run_id}' tracked.",
            resp_json=resp_inference_data,
            ID_path=SOURCE_FILE,
            ID_class=Analyser.__name__,
            ID_function=Analyser._poll_for_stats.__name__
        )

        # Extract the relevant expt-run results
        expt_run_key = replicate_combination_key(self.expt_id, self.run_id)
        metadata = resp_inference_data['results'][expt_run_key]

        # Filter by selected meta-datasets
        filtered_statistics = {
            meta: stats 
            for meta, stats in metadata.items()
            if meta in self.metas
        }

        return {participant_id: filtered_statistics}


    async def _collect_all_stats(self, grid: List[Dict[str, Any]]) -> dict:
        """ Asynchronous function to submit inference data to registered
            participant servers in return for remote performance statistics

        Args:
            grid (list(dict))): Registry of participants' node information
        Returns:
            All participants' statistics (dict)
        """

        ###########################
        # Implementation Footnote #
        ###########################

        # [Causes]
        # Node info retrieved is asynmetric as compared to input inferences.

        # [Problems]
        # Zipping causes misallocation of dataset inferences to the wrong
        # servers, resulting in a "ValueError: Found input variables with 
        # inconsistent numbers of samples: [N1, N2]", since prediction counts 
        # obtained are not aligned with label counts in the remote servers 

        # [Solution]
        # Pair things up via manual matching

        mapped_pairs = [
            (record, inferences)
            for worker_id, inferences in self.inferences.items()
            for record in grid
            if self.parse_syft_info(record).get('id') == worker_id
        ]

        logging.debug(f"Mapped pairs: {mapped_pairs}")

        all_statistics = {}
        for future in asyncio.as_completed(
            map(lambda args: self._poll_for_stats(*args), mapped_pairs)
        ):
            result = await future
            all_statistics.update(result)

        return all_statistics

    ##################
    # Core Functions #
    ##################

    def infer(self, grid: List[Dict[str, Any]]) -> dict:
        """ Wrapper function for triggering asychroneous remote inferencing of
            participant nodes

        Args:
            grid (list(dict))): Registry of participants' node information
        Returns:
            All participants' statistics (dict)
        """
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        try:
            all_stats = loop.run_until_complete(
                self._collect_all_stats(grid=grid)
            )
        finally:
            loop.close()

        return all_stats



####################################
# MLFLow logging class - MLFlogger #
####################################

class MLFlogger:
    """ 
    Wrapper class around MLFlow to faciliate experiment & run registrations in
    the REST-RPC setting, where statistical logging is performed during 
    post-mortem analysis
    """
    def __init__(self, db_path: str = db_path, remote_uri: str = None):

        # Private attributes
        self.__rpc_formatter = RPCFormatter()
        self.__expt_records = ExperimentRecords(db_path=db_path)
        self.__run_records = RunRecords(db_path=db_path)
        self.__model_records = ModelRecords(db_path=db_path)
        self.__connector = "_>_"
        
        # Public attributes
        self.mlf_records = MLFRecords(db_path=db_path)
        self.db_path = db_path
        self.remote_uri = remote_uri

    ###########
    # Helpers #
    ###########

    def _generate_expt_name(
        self, 
        collab_id: str, 
        project_id: str, 
        expt_id: str
    ) -> str:
        """ Abstracts creation of a unique key to identify an MLFlow experiment. 
            This collapses the payload hierarchy so that only one MLFlow instance
            is required to host all runs registered under collaboration(s)  

        Args:
            collab_id (str): ID of collaboration
            project_id (str): ID of project
            expt_id (str): ID of experiment
        Returns:
            Experiment name (str)
        """
        return self.__connector.join([collab_id, project_id, expt_id])


    def _generate_run_name(
        self,         
        collab_id: str, 
        project_id: str, 
        expt_id: str,
        run_id: str
    ) -> str:
        """ Abstracts creation of a unique key to identify an MLFlow run.

        Args:
            collab_id (str): ID of collaboration
            project_id (str): ID of project
            expt_id (str): ID of experiment
            run_id (str): ID of run
        Returns:
            Run name (str)
        """
        return run_id   # keeping it simple

    
    def _generate_record_name(
        self,
        collab_id: str, 
        project_id: str, 
        expt_id: str,
        run_id: str = "",
        mlflow_type: str = "experiment"
    ):
        """ Abstracts creation of a unique key to partition experiments & runs
            internally within the context of MLFlow.

        Args:
            collab_id (str): ID of collaboration
            project_id (str): ID of project
            expt_id (str): ID of experiment
            run_id (str): ID of run
            mlflow_type (str): Type of entity to be represented in MLFlow
        Returns:
            Record name (str)
        """
        
        ###########################
        # Implementation Footnote #
        ###########################

        # [Causes]
        # In MLFlow, only experiments and runs are mappable/usable to the 
        # Synergos hierarchy. In order to simplify the deployment process, 
        # collaborations & projects are collapsed in as part of an MLFlow's 
        # experiment, allowing a single instance of MLFlow to host results &
        # analysis of results across the grid.

        # [Problems]
        # The collapsed hierarchy becomes ambigious due to the inability to
        # distinguish between a Synergos experiment and a Synergos run.

        # [Solution]
        # Use expt_id and run_id to create a unique composite key for use when
        # storing mapping details between MLFlow and Synergos.

        if mlflow_type == "experiment":
            return expt_id

        elif mlflow_type == "run":
            return self.__connector.join([expt_id, run_id])

        else:
            raise RuntimeError("Unsupported entity specified!")


    def retrieve_mlflow_experiment(
        self,
        collab_id: str, 
        project_id: str,
        expt_id: str,
    ) -> Dict[str, str]:
        """
        """
        expt_name = self._generate_expt_name(collab_id, project_id, expt_id)
        expt_record_name = self._generate_record_name(
            collab_id, project_id, expt_id,
            mlflow_type="experiment"
        )
        expt_mlflow_record = self.mlf_records.read(
            collaboration=collab_id,
            project=project_id, 
            record=expt_record_name,
            name=expt_name
        )
        return expt_mlflow_record


    def retrieve_mlflow_run(
        self,
        collab_id: str, 
        project_id: str,
        expt_id: str,
        run_id: str
    ) -> Dict[str, str]:
        """
        """
        run_name = self._generate_run_name(collab_id, project_id, expt_id, run_id)
        run_record_name = self._generate_record_name(
            collab_id, project_id, expt_id, run_id,
            mlflow_type="run"
        )
        run_mlflow_record = self.mlf_records.read(
            collaboration=collab_id,
            project=project_id, 
            record=run_record_name,
            name=run_name
        )
        return run_mlflow_record


    def initialise_mlflow_experiment(
        self, 
        collab_id: str, 
        project_id: str,
        expt_id: str,
        tracking_uri: str
    ) -> Dict[str, str]:
        """ Initialises an MLFlow experiment at the specified project URI

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
        Returns:
            MLFlow Experiment configuration (dict)
        """
        mlflow.set_tracking_uri(tracking_uri)

        expt_name = self._generate_expt_name(collab_id, project_id, expt_id)
        expt_record_name = self._generate_record_name(
            collab_id, project_id, expt_id,
            mlflow_type="experiment"
        )
        expt_mlflow_record = self.mlf_records.read(
            collaboration=collab_id,
            project=project_id, 
            record=expt_record_name,
            name=expt_name
        )

        # Check if MLFlow experiment has already been created. Create a 
        # MLFRecord if: 
        # 1) No previous records for this combination key was initialized
        # 2) Specified tracking URI has changed

        if not expt_mlflow_record:
            
            # Initialise MLFlow experiment
            mlflow_id = mlflow.create_experiment(name=expt_name)

            mlflow_details = {
                'collaboration': collab_id,
                'project': project_id,
                'record': expt_record_name,
                'name': expt_name,
                'mlflow_type': 'experiment',
                'mlflow_id': mlflow_id,
                'mlflow_uri': tracking_uri  # Overridable with remote > local
            }
            expt_mlflow_record = self.mlf_records.create(
                collaboration=collab_id,
                project=project_id,
                record=expt_record_name,
                name=expt_name, 
                details=mlflow_details
            )

        # Update URI if necessary
        if expt_mlflow_record.get('mlflow_uri') != tracking_uri:
            expt_mlflow_record = self.mlf_records.update(
                collaboration=collab_id,
                project=project_id,
                record=expt_record_name,
                name=expt_name, 
                updates={'mlflow_uri': tracking_uri}
            )

        stripped_expt_mlflow_details = self.__rpc_formatter.strip_keys(
            record=expt_mlflow_record
        )
        return stripped_expt_mlflow_details


    def initialise_mlflow_run(
        self, 
        collab_id: str, 
        project_id: str, 
        expt_id: str, 
        run_id: str
    ) -> Dict[str, str]:
        """ Initialises a MLFLow run under a specified experiment of a project.
            Initial run hyperparameters will be logged, and MLFlow run id will
            be stored for subsequent analysis.

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
            run_id (str): REST-RPC ID of specified run
        """
        # Retrieve parent MLFlow experiment metadata (if available)
        expt_mlflow_record = self.retrieve_mlflow_experiment(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id
        )
        expt_mlflow_id = expt_mlflow_record.get('mlflow_id')
        expt_mlflow_uri = expt_mlflow_record.get('mlflow_uri')

        if not expt_mlflow_id:
            logging.error(
                "MLFlow experiment has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.initialise_mlflow_run.__name__
            )
            raise RuntimeError("MLFlow experiment has not been initialised!")

        run_name = self._generate_run_name(collab_id, project_id, expt_id, run_id)
        run_record_name = self._generate_record_name(
            collab_id, project_id, expt_id, run_id,
            mlflow_type="run"
        )
        run_mlflow_record = self.mlf_records.read(
            collaboration=collab_id,
            project=project_id, 
            record=run_record_name,
            name=run_name
        )
        if not run_mlflow_record:

            with mlflow.start_run(
                experiment_id=expt_mlflow_id, 
                run_name=run_name
            ) as mlf_run:

                # Retrieve run details from database
                run_details = self.__run_records.read(
                    collab_id=collab_id, 
                    project_id=project_id, 
                    expt_id=expt_id, 
                    run_id=run_id
                )
                stripped_run_details = self.__rpc_formatter.strip_keys(
                    record=run_details,
                    concise=True
                )

                mlflow.log_params(stripped_run_details)

                # Save the MLFlow ID mapping
                run_mlflow_id = mlf_run.info.run_id
                run_mlflow_details = {
                    'collaboration': collab_id,
                    'project': project_id,
                    'record': run_record_name,
                    'name': run_name,
                    'mlflow_type': 'run',
                    'mlflow_id': run_mlflow_id,
                    'mlflow_uri': expt_mlflow_record['mlflow_uri'] # same as expt
                }
                run_mlflow_record = self.mlf_records.create(
                    collaboration=collab_id,
                    project=project_id,
                    record=run_record_name,
                    name=run_name, 
                    details=run_mlflow_details
                )

        if run_mlflow_record.get('mlflow_uri') != expt_mlflow_uri:
            run_mlflow_record = self.mlf_records.update(
                collaboration=collab_id,
                project=project_id,
                record=run_record_name,
                name=run_name, 
                updates={'mlflow_uri': expt_mlflow_record['mlflow_uri']}
            )

        stripped_run_mlflow_details = self.__rpc_formatter.strip_keys(
            record=run_mlflow_record
        )
        return stripped_run_mlflow_details       


    def log_losses(
        self, 
        collab_id: str, 
        project_id: str, 
        expt_id: str, 
        run_id: str
    ):
        """ Registers all cached losses, be it global or local, obtained from
            federated training into MLFlow.

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
            run_id (str): REST-RPC ID of specified run
        Returns:
            Stripped metadata (dict)
        """
        # Retrieve parent MLFlow experiment metadata (if available)
        expt_mlflow_record = self.retrieve_mlflow_experiment(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id
        )
        expt_mlflow_id = expt_mlflow_record.get('mlflow_id')

        if not expt_mlflow_id:
            logging.error(
                "MLFlow experiment has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow experiment has not been initialised!")

        # Search for run session to update entry, not create a new one
        run_mlflow_record = self.retrieve_mlflow_run(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id, 
            run_id=run_id
        )
        run_mlflow_id = run_mlflow_record.get('mlflow_id')

        if not run_mlflow_id:
            logging.error(
                "MLFlow run has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow run has not been initialised!")

        # Retrieve all model metadata from storage
        model_metadata = self.__model_records.read(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id, 
            run_id=run_id
        )
        stripped_metadata = self.__rpc_formatter.strip_keys(
            record=model_metadata, 
            concise=True
        ) if model_metadata else model_metadata

        if stripped_metadata:

            with mlflow.start_run(
                experiment_id=expt_mlflow_id, 
                run_id=run_mlflow_id # continue existing run
            ) as mlf_run:

                # Extract loss histories
                for m_type, metadata in stripped_metadata.items():

                    loss_history_path = metadata['loss_history']
                    origin = metadata['origin']

                    with open(loss_history_path, 'r') as lh:
                        loss_history = json.load(lh)

                    if m_type == 'global':
                        for meta, losses in loss_history.items():
                            for round_idx, loss, in losses.items():
                                mlflow.log_metric(
                                    key=f"global_{meta}_loss", 
                                    value=loss, 
                                    step=int(round_idx)
                                )

                    else:
                        for round_idx, loss, in losses.items():
                            mlflow.log_metric(
                                key=f"{origin}_local_loss", 
                                value=loss, 
                                step=int(round_idx)
                            )

        return stripped_metadata

    
    def log_model_performance(
        self, 
        collab_id: str,
        project_id: str, 
        expt_id: str, 
        run_id: str,
        statistics: dict
    ) -> dict:
        """ Using all cached model checkpoints, log the performance statistics 
            of models, be it global or local, to MLFlow at round (for global) or
            epoch (for local) level.

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
            run_id (str): REST-RPC ID of specified run
            statistics (dict): Inference statistics polled from workers
        Returns:
            MLFLow run details (dict)
        """
        # Retrieve parent MLFlow experiment metadata (if available)
        expt_mlflow_record = self.retrieve_mlflow_experiment(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id
        )
        expt_mlflow_id = expt_mlflow_record.get('mlflow_id')

        if not expt_mlflow_id:
            logging.error(
                "MLFlow experiment has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow experiment has not been initialised!")

        # Search for run session to update entry, not create a new one
        run_mlflow_record = self.retrieve_mlflow_run(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id, 
            run_id=run_id
        )
        run_mlflow_id = run_mlflow_record.get('mlflow_id')

        if not run_mlflow_id:
            logging.error(
                "MLFlow run has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow run has not been initialised!")

        with mlflow.start_run(
            experiment_id=expt_mlflow_id, 
            run_id=run_mlflow_id
        ) as mlf_run:

            # Store output metadata into database
            for _, inference_stats in statistics.items():
                for _, meta_stats in inference_stats.items():

                    # Log statistics to MLFlow for analysis
                    stats = meta_stats.get('statistics', {})
                    for stat_type, stat_value in stats.items():

                        if isinstance(stat_value, list):
                            for val_idx, value in enumerate(stat_value):
                                mlflow.log_metric(
                                    key=f"{stat_type}_class_{val_idx}", 
                                    value=value, 
                                    step=int(val_idx+1)
                                )

                        else:
                            mlflow.log_metric(key=stat_type, value=stat_value)

        stripped_mlflow_run_details = self.__rpc_formatter.strip_keys(
            run_mlflow_record, 
            concise=True
        )
        return stripped_mlflow_run_details


    def log_artifacts(
        self, 
        collab_id: str,
        project_id: str, 
        expt_id: str, 
        run_id: str
    ) -> str:
        """ Log all artifacts produced (i.e. local & global models, losses etc.)
            during the specified federated run.

        Args:
            project_id (str): REST-RPC ID of specified project
            expt_id (str): REST-RPC ID of specified experiment
            run_id (str): REST-RPC ID of specified run
            statistics (dict): Inference statistics polled from workers
        Returns:
            MLFLow run details (dict)
        """
        # Retrieve parent MLFlow experiment metadata (if available)
        expt_mlflow_record = self.retrieve_mlflow_experiment(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id
        )
        expt_mlflow_id = expt_mlflow_record.get('mlflow_id')

        if not expt_mlflow_id:
            logging.error(
                "MLFlow experiment has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow experiment has not been initialised!")

        # Search for run session to update entry, not create a new one
        run_mlflow_record = self.retrieve_mlflow_run(
            collab_id=collab_id, 
            project_id=project_id, 
            expt_id=expt_id, 
            run_id=run_id
        )
        run_mlflow_id = run_mlflow_record.get('mlflow_id')

        if not run_mlflow_id:
            logging.error(
                "MLFlow run has not been initialised!",
                ID_path=SOURCE_FILE,
                ID_class=MLFlogger.__name__,
                ID_function=MLFlogger.log_losses.__name__
            )
            raise RuntimeError("MLFlow run has not been initialised!")
       
        with mlflow.start_run(
            experiment_id=expt_mlflow_id, 
            run_id=run_mlflow_id # continue existing run
        ) as mlf_run:

            results_dir = os.path.join(out_dir, collab_id, project_id, expt_id, run_id)           
            mlflow.log_artifacts(local_dir=results_dir)

        return results_dir

    ##################
    # Core Functions #
    ##################

    def log(self, accumulations: dict) -> List[str]:
        """ Wrapper function that processes statistics accumulated from 
            inference.

        Args:
            accumulations (dict): Accumulated statistics from inferring
                different project-expt-run combinations
        Returns:
            List of MLFlow run IDs from all runs executed (list(str))
        """

        ###########################
        # Implementation Footnote #
        ###########################

        # [Causes]
        # Depending on where your MLFlow logs are, the accessibility of the
        # resource changes.

        # [Problems]
        # As Synergos MLOps is an optional component, the orchestrator may or
        # may not install it. Also, depending on the location of deployment
        # (i.e. local filesystem vs remote tracking server), connection URI
        # may vary.

        # [Solution]
        # By default, ALWAYS generate local MLFlow caches. If a remote tracking
        # URI was specified, then attempt to send logs over. That way, clients
        # will always have the option of choosing which way they would to serve
        # their MLOps information

        has_remote_tracking = self.remote_uri is not None
        tracking_uris = (
            [mlflow_dir, self.remote_uri] 
            if has_remote_tracking 
            else [mlflow_dir]
        ) 

        jobs_ran = []
        for tracking_uri in tracking_uris:
            for combination_key, statistics in accumulations.items():

                collab_id, project_id, expt_id, run_id = combination_key

                self.initialise_mlflow_experiment(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    tracking_uri=tracking_uri
                )
                run_mlflow_details = self.initialise_mlflow_run(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    run_id=run_id
                )
                self.log_losses(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    run_id=run_id
                )
                self.log_model_performance(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    run_id=run_id,
                    statistics=statistics
                )
                self.log_artifacts(
                    collab_id=collab_id,
                    project_id=project_id,
                    expt_id=expt_id,
                    run_id=run_id
                )

                if run_mlflow_details['mlflow_id'] not in jobs_ran:
                    jobs_ran.append(run_mlflow_details['mlflow_id'])
        
        return jobs_ran
#!/usr/bin/env python #################### # Required Modules # #################### # Generic/Built-in # Libs # Custom ################## # Configurations # ################## #################### # Helper Functions # #################### ############################################ # Inference Orchestration class - Analyser # ############################################ Takes in a list of minibatch IDs and sends them to worker nodes. Workers     will use these IDs to reconstruct their aggregated test datasets with      prediction labels mapped appropriately.      Attributes:         inferences (dict(str, dict(str, dict(str, th.tensor))) ########### # Helpers # ########### Parses a registration record for participant metadata, before             submitting minibatch IDs of inference objects to corresponding              worker node's REST-RPC service for calculating descriptive             statistics and prediction exports          Args:             reg_record (tinydb.database.Document): Participant-project details             inferences (dict): List of dicts containing inference object IDs         Returns:             Statistics (dict) # Construct destination url for interfacing with worker REST-RPC # Trigger remote inference by posting alignments & ID mappings to # `Predict` route in worker # Extract the relevant expt-run results # Filter by selected meta-datasets Asynchronous function to submit inference data to registered             participant servers in return for remote performance statistics          Args:             grid (list(dict))): Registry of participants' node information         Returns:             All participants' statistics (dict) ########################### # Implementation Footnote # ########################### # [Causes] # Node info retrieved is asynmetric as compared to input inferences. # [Problems] # Zipping causes misallocation of dataset inferences to the wrong # servers, resulting in a "ValueError: Found input variables with # inconsistent numbers of samples: [N1, N2]", since prediction counts # obtained are not aligned with label counts in the remote servers # [Solution] # Pair things up via manual matching ################## # Core Functions # ################## Wrapper function for triggering asychroneous remote inferencing of             participant nodes          Args:             grid (list(dict))): Registry of participants' node information         Returns:             All participants' statistics (dict) #################################### # MLFLow logging class - MLFlogger # #################################### Wrapper class around MLFlow to faciliate experiment & run registrations in     the REST-RPC setting, where statistical logging is performed during      post-mortem analysis # Private attributes # Public attributes ########### # Helpers # ########### Abstracts creation of a unique key to identify an MLFlow experiment.              This collapses the payload hierarchy so that only one MLFlow instance             is required to host all runs registered under collaboration(s)            Args:             collab_id (str): ID of collaboration             project_id (str): ID of project             expt_id (str): ID of experiment         Returns:             Experiment name (str) Abstracts creation of a unique key to identify an MLFlow run.          Args:             collab_id (str): ID of collaboration             project_id (str): ID of project             expt_id (str): ID of experiment             run_id (str): ID of run         Returns:             Run name (str) # keeping it simple Abstracts creation of a unique key to partition experiments & runs             internally within the context of MLFlow.          Args:             collab_id (str): ID of collaboration             project_id (str): ID of project             expt_id (str): ID of experiment             run_id (str): ID of run             mlflow_type (str): Type of entity to be represented in MLFlow         Returns:             Record name (str) ########################### # Implementation Footnote # ########################### # [Causes] # In MLFlow, only experiments and runs are mappable/usable to the # Synergos hierarchy. In order to simplify the deployment process, # collaborations & projects are collapsed in as part of an MLFlow's # experiment, allowing a single instance of MLFlow to host results & # analysis of results across the grid. # [Problems] # The collapsed hierarchy becomes ambigious due to the inability to # distinguish between a Synergos experiment and a Synergos run. # [Solution] # Use expt_id and run_id to create a unique composite key for use when # storing mapping details between MLFlow and Synergos. Initialises an MLFlow experiment at the specified project URI          Args:             project_id (str): REST-RPC ID of specified project             expt_id (str): REST-RPC ID of specified experiment         Returns:             MLFlow Experiment configuration (dict) # Check if MLFlow experiment has already been created. Create a # MLFRecord if: # 1) No previous records for this combination key was initialized # 2) Specified tracking URI has changed # Initialise MLFlow experiment # Overridable with remote > local # Update URI if necessary Initialises a MLFLow run under a specified experiment of a project.             Initial run hyperparameters will be logged, and MLFlow run id will             be stored for subsequent analysis.          Args:             project_id (str): REST-RPC ID of specified project             expt_id (str): REST-RPC ID of specified experiment             run_id (str): REST-RPC ID of specified run # Retrieve parent MLFlow experiment metadata (if available) # Retrieve run details from database # Save the MLFlow ID mapping # same as expt Registers all cached losses, be it global or local, obtained from             federated training into MLFlow.          Args:             project_id (str): REST-RPC ID of specified project             expt_id (str): REST-RPC ID of specified experiment             run_id (str): REST-RPC ID of specified run         Returns:             Stripped metadata (dict) # Retrieve parent MLFlow experiment metadata (if available) # Search for run session to update entry, not create a new one # Retrieve all model metadata from storage # continue existing run # Extract loss histories Using all cached model checkpoints, log the performance statistics              of models, be it global or local, to MLFlow at round (for global) or             epoch (for local) level.          Args:             project_id (str): REST-RPC ID of specified project             expt_id (str): REST-RPC ID of specified experiment             run_id (str): REST-RPC ID of specified run             statistics (dict): Inference statistics polled from workers         Returns:             MLFLow run details (dict) # Retrieve parent MLFlow experiment metadata (if available) # Search for run session to update entry, not create a new one # Store output metadata into database # Log statistics to MLFlow for analysis Log all artifacts produced (i.e. local & global models, losses etc.)             during the specified federated run.          Args:             project_id (str): REST-RPC ID of specified project             expt_id (str): REST-RPC ID of specified experiment             run_id (str): REST-RPC ID of specified run             statistics (dict): Inference statistics polled from workers         Returns:             MLFLow run details (dict) # Retrieve parent MLFlow experiment metadata (if available) # Search for run session to update entry, not create a new one # continue existing run ################## # Core Functions # ################## Wrapper function that processes statistics accumulated from              inference.          Args:             accumulations (dict): Accumulated statistics from inferring                 different project-expt-run combinations         Returns:             List of MLFlow run IDs from all runs executed (list(str)) ########################### # Implementation Footnote # ########################### # [Causes] # Depending on where your MLFlow logs are, the accessibility of the # resource changes. # [Problems] # As Synergos MLOps is an optional component, the orchestrator may or # may not install it. Also, depending on the location of deployment # (i.e. local filesystem vs remote tracking server), connection URI # may vary. # [Solution] # By default, ALWAYS generate local MLFlow caches. If a remote tracking # URI was specified, then attempt to send logs over. That way, clients # will always have the option of choosing which way they would to serve # their MLOps information