backend/aws_handler.py

import logging
import boto3
from botocore.exceptions import ClientError
import os
from datetime import datetime
import torch


def upload_file(file_name, bucket="fishcounting", object_name=None):
    """Upload a file to an S3 bucket

    :param file_name: File to upload
    :param bucket: Bucket to upload to
    :param object_name: S3 object name. If not specified then file_name is used
    :return: True if file was uploaded, else False
    """

    # If S3 object_name was not specified, use file_name
    if object_name is None:
        object_name = os.path.basename(file_name)

    if (not 'AAK_ID' in os.environ) or (not 'ASAK' in os.environ):
        print('AWS keys not specified. Cancelling sync')
        return False

    # Upload the file
    s3_client = boto3.client(
        's3',
        aws_access_key_id=os.environ['AAK_ID'],
        aws_secret_access_key=os.environ['ASAK']
    )
    try:
        response = s3_client.upload_file(file_name, bucket, object_name)
    except ClientError as e:
        logging.error(e)
        return False
    return True

def ping_server(aris_name, state):
    """Upload a notification file to AWS

    :param aris_name: Name of the aris file uploaded
    :return: True if file was uploaded, else False
    """

    file_name = 'tmp/notification.txt'
    os.makedirs('tmp', exist_ok=True)
    with open(file_name, 'w') as f:
        output = f"time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        output += f"filename: {aris_name}\n"
        output += f"app version: {state['version']}\n"
        output += f"hardware: {torch.cuda.get_device_name() if torch.cuda.is_available() else 'CPU'}\n"
        if 'CPU_CORES' in os.environ:
            output += f"location: HuggingFace\n"
            output += f"nbr cpu cores: {os.environ['CPU_CORES']}\n"
            output += f"memory: {os.environ['MEMORY']}\n"
        else:
            output += f"location: local install"
        f.write(output)
 
    # If S3 object_name was not specified, use file_name

    if (not 'AAK_ID' in os.environ) or (not 'ASAK' in os.environ):
        print('AWS keys not specified. Cancelling sync')
        return False

    # Upload the file
    s3_client = boto3.client(
        's3',
        aws_access_key_id=os.environ['AAK_ID'],
        aws_secret_access_key=os.environ['ASAK']
    )
    try:
        response = s3_client.upload_file(file_name, "fishcounting", "webapp_uploads/notifications/" + str(int(datetime.now().timestamp())) + ".txt")
    except ClientError as e:
        logging.error(e)
        return False
    return True