From 80e58761e3f10b52a54acc8337a87e06b8be5614 Mon Sep 17 00:00:00 2001 From: Shay Erlichmen Date: Mon, 14 Oct 2024 00:28:29 +0300 Subject: [PATCH] fix: handle error handling --- skyplane/api/usage.py | 68 +++++++++++++++++++++++++++++++------------ 1 file changed, 49 insertions(+), 19 deletions(-) diff --git a/skyplane/api/usage.py b/skyplane/api/usage.py index f939dcae..61f5990f 100644 --- a/skyplane/api/usage.py +++ b/skyplane/api/usage.py @@ -1,31 +1,29 @@ +import configparser import datetime import json import os import sys import time import uuid -import configparser from dataclasses import asdict, dataclass from enum import Enum, auto from pathlib import Path +from typing import Optional, Dict, List import requests from rich import print as rprint -from typing import Optional, Dict, List import skyplane -from skyplane.utils.definitions import tmp_log_dir from skyplane.config import _map_type from skyplane.config_paths import config_path, cloud_config, host_uuid_path from skyplane.utils import logger, imports +from skyplane.utils.definitions import tmp_log_dir SCHEMA_VERSION = "0.2" LOKI_URL = "http://34.212.234.105:9090/loki/api/v1/push" USAGE_STATS_ENABLED_ENV_VAR = "SKYPLANE_USAGE_STATS_ENABLED" USAGE_STATS_FILE = "usage_stats.json" -USAGE_STATS_ENABLED_MESSAGE = ( - "[bright_black]To disable performance logging info: https://skyplane.org/en/latest/performance_stats_collection.html[/bright_black]" -) +USAGE_STATS_ENABLED_MESSAGE = "[bright_black]To disable performance logging info: https://skyplane.org/en/latest/performance_stats_collection.html[/bright_black]" USAGE_STATS_DISABLED_RECONFIRMATION_MESSAGE = ( "[green][bold]We are an academic research group working to improve inter-cloud network performance.[/bold] " "You can inspect what we share in the /tmp/skyplane/metrics directory. " @@ -38,9 +36,7 @@ "\n * Total runtime and the aggregated transfer speed in Gbps" "\n * Error message if the transfer fails" ) -USAGE_STATS_REENABLE_MESSAGE = ( - "[yellow][bold]If you want to re-enable usage statistics, run `skyplane config set usage_stats true`.[/bold][/yellow]" -) +USAGE_STATS_REENABLE_MESSAGE = "[yellow][bold]If you want to re-enable usage statistics, run `skyplane config set usage_stats true`.[/bold][/yellow]" USAGE_STATS_REENABLED_MESSAGE = ( "[green][bold]Thank you for your support of open-source research![/bold][/green]" "\nIf you want to disable usage statistics, run `skyplane config set usage_stats false`." @@ -219,18 +215,26 @@ def usage_stats_status(cls) -> UsageStatsStatus: def set_usage_stats_via_config(typer, cls, value, config): current_status = cls.usage_stats_status() if current_status is UsageStatsStatus.DISABLED_EXPLICITLY: - if (isinstance(value, bool) and not value) or (isinstance(value, str) and not _map_type(value, bool)): + if (isinstance(value, bool) and not value) or ( + isinstance(value, str) and not _map_type(value, bool) + ): rprint("Usage stats collection is already disabled.") rprint(USAGE_STATS_REENABLE_MESSAGE) return elif current_status is UsageStatsStatus.ENABLED_EXPLICITLY: - if (isinstance(value, bool) and value) or (isinstance(value, str) and _map_type(value, bool)): + if (isinstance(value, bool) and value) or ( + isinstance(value, str) and _map_type(value, bool) + ): rprint("Usage stats collection is already enabled.") rprint(USAGE_STATS_REENABLED_MESSAGE) return - if (isinstance(value, bool) and not value) or (isinstance(value, str) and not _map_type(value, bool)): - prompt = "Would you still like to opt out of sharing anonymous usage metrics?" + if (isinstance(value, bool) and not value) or ( + isinstance(value, str) and not _map_type(value, bool) + ): + prompt = ( + "Would you still like to opt out of sharing anonymous usage metrics?" + ) rprint(USAGE_STATS_DISABLED_RECONFIRMATION_MESSAGE + "\n") answer = typer.confirm(prompt, default=False) if not answer: @@ -242,7 +246,9 @@ def set_usage_stats_via_config(typer, cls, value, config): try: config.set_flag("usage_stats", value) except Exception as e: - raise Exception("Failed to enable/disable by writing to" f"{config_path}") from e + raise Exception( + "Failed to enable/disable by writing to" f"{config_path}" + ) from e if config.get_flag("usage_stats"): rprint(USAGE_STATS_REENABLED_MESSAGE) @@ -265,6 +271,9 @@ def make_stat( if dest_region_tags is None: dest_provider, dest_region = None, None else: + if isinstance(dest_region_tags, str): + dest_region_tags = [dest_region_tags] + dest_providers = [tag.split(":")[0] for tag in dest_region_tags] dest_regions = [tag.split(":")[1] for tag in dest_region_tags] @@ -279,7 +288,11 @@ def make_stat( source_cloud_provider=src_provider, destination_cloud_providers=dest_providers, os=sys.platform, - session_start_timestamp_ms=session_start_timestamp_ms if session_start_timestamp_ms else int(time.time() * 1000), + session_start_timestamp_ms=( + session_start_timestamp_ms + if session_start_timestamp_ms + else int(time.time() * 1000) + ), arguments_dict=arguments_dict, transfer_stats=transfer_stats, ) @@ -316,12 +329,18 @@ def make_error( source_cloud_provider=src_provider, destination_cloud_providers=dest_providers, # TODO: FIX THIS os=sys.platform, - session_start_timestamp_ms=session_start_timestamp_ms if session_start_timestamp_ms else int(time.time() * 1000), + session_start_timestamp_ms=( + session_start_timestamp_ms + if session_start_timestamp_ms + else int(time.time() * 1000) + ), arguments_dict=arguments_dict, error_dict=error_dict, ) - def write_usage_data(self, data: UsageStatsToReport, dir_path: Optional[Path] = None): + def write_usage_data( + self, data: UsageStatsToReport, dir_path: Optional[Path] = None + ): """Write the usage data to the directory. :param data: data to report @@ -342,7 +361,9 @@ def write_usage_data(self, data: UsageStatsToReport, dir_path: Optional[Path] = json_file.write(json.dumps(asdict(data))) return destination - def report_usage_data(self, type: str, data: UsageStatsToReport, path: Path) -> None: + def report_usage_data( + self, type: str, data: UsageStatsToReport, path: Path + ) -> None: """Report the usage data to the usage server. :param type: the type of usage stats to report @@ -356,7 +377,16 @@ def report_usage_data(self, type: str, data: UsageStatsToReport, path: Path) -> prom_labels = {"type": type, "environment": "api"} headers = {"Content-type": "application/json"} data.sent_time_ms = int(time.time() * 1000) - payload = {"streams": [{"stream": prom_labels, "values": [[str(_get_current_timestamp_ns()), json.dumps(asdict(data))]]}]} + payload = { + "streams": [ + { + "stream": prom_labels, + "values": [ + [str(_get_current_timestamp_ns()), json.dumps(asdict(data))] + ], + } + ] + } payload = json.dumps(payload) r = requests.post(LOKI_URL, headers=headers, data=payload, timeout=0.5)