timeseries-experiments/df_api.py

import os
import time
import copy
import json
import uuid
import boto3
import requests

from tqdm import tqdm
from typing import List, Tuple, Optional


#AWS
AWS_ACCESS_KEY_ID = 'AKIAI6PFXVX6CBIHIJIA'
AWS_SECRET_ACCESS_KEY = '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'

#Anastasia 
CLIENT_UUID = 'fed091cb-754c-42fc-8d35-cd28601dfc75'
PROJECT_UUID = '79d14f32-02b9-45fb-9392-284fb7c6a251'
APPLICATION_UUID = 'f28bcce0-c521-41ae-9db7-f985de54272d'
DATASOURCE_TYPE = 'T001'    #tipo transaccional

#APIs
INTERNAL_DEV_API = 'http://int-srv-dev.anastasia.ai/'
X_API_KEY = 'J48NWuVnLcFmavDH37wUvMUVKGd8Edda6VhVmDhMJj7xo'
DATASOURCES_API_KEY = 'MM9gvWjG3Z522UfNkEu9ZbLF3r6e9jqBF33PY4z3usHV7mtE7w'

#Datasets
OPENDATA_COLUMNS = [
    "651c8079-a9d7-4ec9-8a5a-8ecb2375fdc6", "4492e66b-b569-4bd9-885c-ff504a64021e",
    "cd1bf19b-9f5b-4244-bdb1-f4c892fa02e1", "c546adf2-ff80-4965-9ec1-3982c621763e",
    "b0ae454e-bf5b-4e18-b35f-b416d56ec6da", "5d63e743-e4f9-432b-8fed-dfe2d0e9811d",
    "238c38eb-dc6c-4690-b89a-2693e05cfa1d", "ab3c92f1-32fe-463d-bb3b-15f79f8ec466",
    "3152982c-7809-4801-b63d-aba822f49798", "41a7b328-b955-4c2a-81ad-78100074942e",
    "a4c5077e-2750-4f32-8e27-0f4be9d61470", "f62385f1-2604-4e5c-adcf-af374e0a727f",
    "3eb5f163-ff9a-4e86-876b-3c41d4f9298a", "8874a1eb-89cb-4489-b552-219ffef0d5c4",
    "a8d989c5-2607-4af3-a3bb-acd590c30f74", "328bc204-99ad-4600-9aa1-588326cc7778",
    "4165f7cd-0760-4f2e-a262-fa74aa4a4829", "24aa83be-b4d6-40d1-b429-69977c4e3282",
    "98a06172-0637-412e-beab-4ae1e7c1b05e", "2a85dbcd-b13f-4a58-b584-050248b16635",
    "5ff95e31-40c1-4579-8a7e-2b893c8eba89", "7aeba544-7b1c-4baa-a0dc-beaf1145b09c",
    "433067ce-8be6-43cb-a68c-925690cb00be"
]


class DemandForecasting:

    def __init__(self, verbose: bool = True) -> None:
        self.verbose = verbose

    def upload_datasource(self, name: str, columns: List[dict], files_path: str) -> str:
        """Crea un datasource y carga archivos hacia el mismo

        Parameters
        ----------

        name : str
            Nombre que tendra el datasource

        columns : List[dict]
            Descripción de estructura de columnas

        files_path : str
            Ruta en la cual se ubican los archivos a subir

        Returns
        -------

        datasource_uuid : str
            Uuid de identificación del datasource creado con los archivos cargados
        """
        datasource_columns = self._translate_columns_types(columns)
        datasource_uuid = self._create_datasource(name, datasource_columns)['code']
        datasource_info = self._get_datasource_info(datasource_uuid)
        s3_path = datasource_info['storageURL']
        datasets_uuids = self._upload_multi_files(files_path, datasource_uuid,s3_path)

        return datasource_uuid

    def run_job(self, prediction_params: dict, datasource_params: dict, extras: Optional[dict] = None, wait_until_finished: bool = True, timeout: int = 2*60*60) -> str:
        """Invocar a la API de jobs ejecutando la solucion Demand Forecasting

        Parameters
        ----------

        prediction_params : dict
            Descripcion de parametros relacionados a la prediccion

        datasource_params : dict
            Descripcion de estructura de columnas para asociar hacia el predictor

        extras : dict
            Parametros extras para modificar comportamiento de Predictor

        wait_until_finished : bool, default=True
            Bloquea la ejecucion de codigo posterior a la espera de que el job en la nube termine con un status de parada

        timeout : int, default=7200 (unidades en segundos)
            Tolerancia maxima de espera para esperar un status de finalizacion (siempre que la variable wait_until_finished sea gatillada)

        Returns
        -------

        job_uuid : str
            Uuid de identificación del job ejecutado
        """
        datasource_uuid = datasource_params["datasource_uuid"]
        use_opendata = datasource_params.get("use_opendata", None)
        datasource_info = self._get_datasource_info(datasource_uuid)
        columns_description = datasource_info["columns"]


        for column in columns_description:
            if column["originalName"] == datasource_params["date_column"]:
                date_column = column["mapName"]
                break

        date_column = [{"column":column["mapName"], "dataSourceCode": datasource_uuid} for column in columns_description if column["originalName"] == datasource_params["date_column"]][0]
        id_columns = [{"column":column["mapName"], "dataSourceCode": datasource_uuid} for column in columns_description if column["originalName"] in datasource_params["id_columns"]]
        target_column = [{"column":column["mapName"], "dataSourceCode": datasource_uuid} for column in columns_description if column["originalName"] == datasource_params["target_column"]][0]

        timeseries_params = {
            "aggregator": "SUM",
            "clientUUID": CLIENT_UUID,
            "projectUUID": PROJECT_UUID,
            "applicationUUID": APPLICATION_UUID,
            "granularity": "MONTH",
            "negativeValues": False,
            "accumulatedValues": False,
            "relationships": [],
            "extraValues": [],
            "eventsDataSourceCode": datasource_uuid,
            "valueColumn": target_column,
            "groupableColumns": id_columns,
            "dateColumn": date_column,
        }
        if use_opendata: timeseries_params["openData"] = OPENDATA_COLUMNS

        job_config = {
            "project_uuid": PROJECT_UUID,
            "client_uuid": CLIENT_UUID,
            "application_uuid": APPLICATION_UUID,
            "params": {
                "core": {},
                "platform": {
                    "ray_params": {
                        "ray_image": "510629978475.dkr.ecr.us-east-1.amazonaws.com/core-ai-library-dev",
                        "head_instance_type": "m5.2xlarge",
                        "worker_instance_type": "m5.2xlarge",
                        "max_workers": 10,
                        "job_submission": {
                            "entrypoint": "python entrypoint/predictor.py"
                        }
                    },
                    "time_series_params": timeseries_params,
                    "prediction_params": prediction_params,
                    "extras": extras
                }
            }
        }
        
        job_type = 'DEMAND_FORECASTING'
        url = self._merge_paths(INTERNAL_DEV_API, f'jobs/api/v1/{job_type}')
        payload = json.dumps(job_config)
        headers = {
            'x-anastasia-api-key': X_API_KEY
        }
        try:
            response = requests.request("POST", url, headers=headers, data=payload)
            if response.status_code != 200:
                print(response.json())
                raise
        except Exception as e:
            print("Error while configuring job")
            print(e.text)
            raise

        job_uuid = response.json()["uuid"]
        if wait_until_finished:
            url = self._merge_paths(INTERNAL_DEV_API, f'/jobs/api/v1/{job_uuid}')
            pbar = tqdm(total=1, desc="Running Demand Forecasting via API in cloud cluster, waiting until job is completed")
            mustend = time.time() + timeout
            while time.time() < mustend:
                try:
                    response = requests.request("GET", url, headers=headers).json()
                except Exception as e:
                    print("Error reading job info")
                    print(e)
                    raise

                if response["status"] == "FINISHED":
                    pbar.update(1)
                    pbar.close()
                    return job_uuid
                elif response["status"] == "FAILED":
                    print(f"Error:\njob_uuid = {job_uuid}\nstatus = {response['status']}")
                    pbar.close()
                    return job_uuid

            print("error: timeout waiting job!")
            pbar.close()
            raise
        else:
            return job_uuid

    def get_job_status(self, job_uuid: str) -> str:
        """Consulta acerca del status de un job

        Parameters
        ----------

        job_uuid : str
            Uuid de identificación del job ejecutado

        Returns
        -------

        status : str
            Estado del job
        """
        headers = {
            'x-anastasia-api-key': X_API_KEY
        }
        url = self._merge_paths(INTERNAL_DEV_API, f'/jobs/api/v1/{job_uuid}')
        try:
            response = requests.request("GET", url, headers=headers).json()
        except Exception as e:
            print("Error reading job info")
            print(e)
            raise

        return response["status"]

    def download_results(self, job_uuid: str, save_path: str) -> None:
        """Descarga los archivos resultantes de un job de Demand Forecasting

        Parameters
        ----------

        job_uuid : str
            Uuid de identificación del job ejecutado

        save_path : str
            Directorio en el cual se descargaran los archivos ["results.csv", "metrics.csv", "item_classes.csv", "predictions_confidence.csv"] 

        """
        headers = {
            'x-anastasia-api-key': X_API_KEY
        }
        url = self._merge_paths(INTERNAL_DEV_API, f'/jobs/api/v1/{job_uuid}')
        try:
            response = requests.request("GET", url, headers=headers).json()
        except Exception as e:
            print("Error reading job info")
            print(e)
            raise

        cloud_path = response["state"]["core"]["dataset"]["output"]["path"]
        file_names = ["results.csv", "metrics.csv", "item_classes.csv", "predictions_confidence.csv"]

        s3_client = boto3.client(
            "s3",
            aws_access_key_id=AWS_ACCESS_KEY_ID,
            aws_secret_access_key=AWS_SECRET_ACCESS_KEY)
        bucket, filekey = self._decompose_s3_route(cloud_path)
        try:
            for file in file_names:
                response = s3_client.download_file(
                    bucket,
                    self._merge_paths(filekey,file),
                    self._merge_paths(save_path,file),
                )
                if self.verbose:
                    print(f"Download successful to {save_path}")
        except Exception as e:
            print("Error reading save_path")
            print(e)
            raise

    def _merge_paths(self, path_1: str, path_2: str):

        if path_1[-1] == '/' and path_2[0] == '/':
            return path_1 + path_2[1:]
        if path_1[-1] != '/' and path_2[0] != '/^':
            return path_1 + '/' + path_2

        return path_1 + path_2

    def _translate_columns_types(self, columns: List[dict]) -> List[dict]:

        code_map = {
            'date':'C001',
            'groupable': 'C002',
            'item': 'C003',
            'value': 'C004',
            'description': 'C005',
            'complementary': 'C006'
        }
        new_columns = []
        for column in columns:
            new_column = copy.deepcopy(column)
            new_column['columnType'] = code_map[column['columnType']]
            new_columns.append(new_column)

        return new_columns

    def _create_datasource(self, name, columns) -> dict:

        if self.verbose: print("Creating datasource...")
        url = self._merge_paths(INTERNAL_DEV_API, 'datasource/api/v2/datasources/')
        payload= json.dumps({
            "clientUUID": CLIENT_UUID,
            "projectUUID": PROJECT_UUID,
            "dataSourceTypeCode": DATASOURCE_TYPE,
            "name": name,
            "columns": columns
        })
        headers = {
            'data_source-anastasia-api-key': DATASOURCES_API_KEY,
            'Content-Type': 'application/json'
        }

        try:
            response = requests.request("POST", url, headers=headers, data=payload)
        except Exception as e:
            print("Error creating datasource")
            print(e.text)
            raise
        if self.verbose:
            print("Datasource created!")

        return response.json()

    def _get_datasource_info(self, datasource_uuid: str) -> dict:

        if self.verbose: print("Getting datasource info...")
        params = 'columns=true&storageLocation=internal&athenaTable=true'
        url = self._merge_paths(INTERNAL_DEV_API,f'datasource/api/v2/datasources/{datasource_uuid}?{params}')
        payload={}
        headers = {
            'data_source-anastasia-api-key': DATASOURCES_API_KEY
        }

        try:
            response = requests.request("GET", url, headers=headers, data=payload)
        except Exception as e:
            print("Error reading datasource")
            print(e.text)
            raise

        return response.json()

    def _upload_multi_files(self, local_datasets_path, datasource_uuid, storageURL) -> dict:
        dataset_uuids = {}
        dataset_names = self._get_datasets_names(local_datasets_path)
        for name in dataset_names:
            dataset_uuid = str(uuid.uuid4())
            self._upload_file(
                output_location=storageURL,
                local_folder=local_datasets_path,
                input_filename = name,
                output_filename=dataset_uuid,
            )
            dataset_uuids[name] = f"{dataset_uuid}.csv"
            file_size = self._get_dataset_size(local_datasets_path,name)
            self._confirm_datasource_loaded(datasource_uuid, name, dataset_uuid, file_size)
        return dataset_uuids

    def _get_datasets_names(self, local_datasets_path: str) -> List[str]:

        datasets = []
        files = os.listdir(local_datasets_path)
        for file in files:
            if file[-4:] == '.csv':
                datasets.append(file)
        assert datasets, 'No data founded!'
        if self.verbose:
            print(f'Total files: {len(datasets)}')

        return datasets

    def _get_dataset_size(self, local_datasets_path: str, name: str) -> int:

        size = os.path.getsize(self._merge_paths(local_datasets_path,name))
        assert size > 0, 'No data founded!'
        if self.verbose:
            print(f'name: {name}, size: {size} bytes')

        return size

    def _decompose_s3_route(self, s3_route: str) -> Tuple[str]:

        s3_route = s3_route.replace("s3://", "").replace("https://", "")
        bucket = s3_route.split("/")[0]
        filekey = "/".join(s3_route.split("/")[1:])

        return bucket, filekey

    def _upload_file(self, output_location, local_folder, input_filename, output_filename) -> bool:

        if self.verbose:
            print("Uploading  file to s3")
        s3_client = boto3.client(
            "s3",
            aws_access_key_id=AWS_ACCESS_KEY_ID,
            aws_secret_access_key=AWS_SECRET_ACCESS_KEY)
        bucket, filekey = self._decompose_s3_route(output_location)
        if output_filename.split(".")[-1] != "csv":
            output_filename = output_filename + ".csv"
        try:
            response = s3_client.upload_file(
                self._merge_paths(local_folder,input_filename),
                bucket,
                filekey + output_filename
            )
            if self.verbose:
                print(f"Upload successful to {output_location}")
                print("-"*50)
        except Exception as e:
            print(e)
            return False
        return True

    def _confirm_datasource_loaded(self, datasource_uuid: str, name: str, dataset_uuid: str, file_size: int):
        if self.verbose: print("Confirming datasource...")
        url = self._merge_paths(INTERNAL_DEV_API,'datasource/api/v2/files')
        payload={
            "name": name,
            "size": file_size,
            "uuid": dataset_uuid,
            "dataSourceCode": datasource_uuid
        }
        headers = {
            'data_source-anastasia-api-key': DATASOURCES_API_KEY
        }
        try:
            response = requests.request("POST", url, headers=headers, data=payload)
        except Exception as e:
            print("Error confirming datasource")
            print(e.text)
            raise
        if self.verbose:
            print("datasource confirmated!")