Spaces:
Runtime error
Runtime error
| import datetime | |
| import json | |
| import os | |
| import pickle | |
| import random | |
| import string | |
| import tempfile | |
| import time | |
| import boto3 | |
| import pandas as pd | |
| from awswrangler import s3 | |
| def generate_random_id(length=8): | |
| # Combine uppercase letters, lowercase letters, and digits | |
| characters = string.ascii_letters + string.digits | |
| # Use random.choices to generate a list of random characters of the specified length | |
| random_id = "".join(random.choices(characters, k=length)) | |
| return random_id | |
| class CustomJSONEncoder(json.JSONEncoder): | |
| """Custom encoder to convert datetime to str if needed.""" | |
| def default(self, obj): | |
| if isinstance(obj, (pd.Timestamp, datetime.date)): | |
| return obj.isoformat() # Convert to ISO 8601 string | |
| return super().default(obj) | |
| def write_content_to_s3_or_local(content, output_path): | |
| extension = output_path.split(".")[-1] # will be pkl or json | |
| if output_path.startswith("s3"): | |
| if extension == "json": | |
| with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: | |
| json.dump(content, temp_file, indent=4) | |
| temp_file_path = temp_file.name | |
| elif extension == "pkl": | |
| with tempfile.NamedTemporaryFile(mode="wb", delete=False) as temp_file: | |
| pickle.dump(content, temp_file) | |
| temp_file_path = temp_file.name | |
| s3.upload(temp_file_path, output_path) | |
| os.remove(temp_file_path) | |
| else: | |
| if extension == "json": | |
| with open(output_path, "wb") as file: | |
| json.dump(content, file, indent=4) | |
| if extension == "pkl": | |
| with open(output_path, "wb") as file: | |
| pickle.dump(content, file) | |
| def read_content_from_s3_or_local(path: str): | |
| extension = path.split(".")[-1] # will be pkl or json | |
| read_from_s3 = path.startswith("s3") | |
| if read_from_s3: | |
| with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: | |
| temp_file_path = temp_file.name | |
| s3.download(path, temp_file_path) | |
| else: | |
| temp_file_path = path | |
| with open(temp_file_path, "rb") as file: | |
| if extension == "json": | |
| data = json.load(file) | |
| elif extension == "pkl": | |
| data = pickle.load(file) | |
| if read_from_s3: | |
| os.remove(temp_file_path) | |
| return data | |
| def start_ec2_instance(region_name: str, instance_id: str): | |
| ec2 = boto3.client("ec2", region_name=region_name) | |
| response = ec2.describe_instances(InstanceIds=[instance_id]) | |
| state = response["Reservations"][0]["Instances"][0]["State"]["Name"] | |
| if state == "stopping": | |
| # Wait until the instance is stopped | |
| while True: | |
| response = ec2.describe_instances(InstanceIds=[instance_id]) | |
| state = response["Reservations"][0]["Instances"][0]["State"]["Name"] | |
| if state == "stopped": | |
| ec2.start_instances(InstanceIds=[instance_id]) | |
| break | |
| time.sleep(5) # Wait for 10 seconds before checking again | |
| else: | |
| ec2.start_instances(InstanceIds=[instance_id]) | |