run_server.py

import argparse
from dataclasses import dataclass, field
import json
import copy
import multiprocessing as mp
import uuid
from datetime import datetime, timedelta
from collections import defaultdict, deque
import io
import zipfile
import queue
import time
import random
import logging

from tensordict import TensorDict
import cv2
from flask import Flask, request, make_response, send_file
from PIL import Image
import torchvision.transforms as T
import numpy as np
import torch as th

from wham.utils import load_model_from_checkpoint, POS_BINS_BOUNDARIES, POS_BINS_MIDDLE

logging.basicConfig(level=logging.INFO)

parser = argparse.ArgumentParser(description="Simple Dreamer")
parser.add_argument("--model", type=str, required=True, help="Path to the model file for the local runs")
parser.add_argument("--debug", action="store_true", help="Enable flask debug mode.")
parser.add_argument("--random_model", action="store_true", help="Use randomly initialized model instead of the provided one")
parser.add_argument("--port", type=int, default=5000)

parser.add_argument("--max_concurrent_jobs", type=int, default=30, help="Maximum number of jobs that can be run concurrently on this server.")
parser.add_argument("--max_dream_steps_per_job", type=int, default=10, help="Maximum number of dream steps each job can request.")
parser.add_argument("--max_job_lifespan", type=int, default=60 * 10, help="Maximum number of seconds we keep run around if not polled.")

parser.add_argument("--image_width", type=int, default=300, help="Width of the image")
parser.add_argument("--image_height", type=int, default=180, help="Height of the image")

parser.add_argument("--max_batch_size", type=int, default=3, help="Maximum batch size for the dreamer workers")

PREDICTION_JSON_FILENAME = "predictions.json"
# Minimum time between times we check when to delete jobs. We do this when adding new jobs.
JOB_CLEANUP_CHECK_RATE = timedelta(seconds=10)

MAX_CANCELLED_ID_QUEUE_SIZE = 100

DEFAULT_SAMPLING_SETTINGS = {
    "temperature": 0.9,
    "top_k": None,
    "top_p": 1.0,
    "max_context_length": 10,
}


def float_or_none(string):
    if string.lower() == "none":
        return None
    return float(string)


def be_image_preprocess(image, target_width, target_height):
    # If target_width and target_height are specified, resize the image.
    if target_width is not None and target_height is not None:
        # Make sure we do not try to resize if the image is already the correct size.
        if image.shape[1] != target_width or image.shape[0] != target_height:
            image = cv2.resize(image, (target_width, target_height))
    return np.transpose(image, (2, 0, 1))


def action_vector_to_be_action_vector(action):
    # Preprocess a BE action vector from 16 numbers with:
    #  12 buttons [0, 1] and 4 stick directions [-1, 1]
    # to discrete actions valid for the token model
    #  12 buttons [0, 1] and 4 stick directions {discrete bin}
    action[-4:] = np.digitize(action[-4:], bins=POS_BINS_BOUNDARIES) - 1
    return action


def be_action_vector_to_action_vector(action):
    # Preprocess a BE action vector into unified space
    for stick_index in range(-4, 0):
        action[stick_index] = POS_BINS_MIDDLE[int(action[stick_index])]
    return action



@dataclass
class DreamJob:
    job_id: str
    sampling_settings: dict
    num_predictions_remaining: int
    num_predictions_done: int
    # (B, T, C, H, W)
    context_images: th.Tensor
    context_actions: th.Tensor
    # Tokens that will replace the context_images if they are provided
    context_tokens: list
    # This will replace the dreamed action if provided.
    # For every step, we remove the first action until exhausted
    actions_to_take: th.Tensor = None


@dataclass
class DreamJobResult:
    job_id: str
    dream_step_index: int
    # (B, 1, C, H, W)
    dreamt_image: th.Tensor
    dreamt_action: th.Tensor
    dreamt_tokens: th.Tensor
    result_creation_time: datetime = field(default_factory=datetime.now)



def setup_and_load_model_be_model(args):
    model = load_model_from_checkpoint(args.model)
    th.set_float32_matmul_precision("high")
    th.backends.cuda.matmul.allow_tf32 = True
    return model


def get_job_batchable_information(job):
    """Return comparable object of job information. Used for batching"""
    context_length = job.context_images.shape[1]
    return (context_length, job.sampling_settings)


def fetch_list_of_batchable_jobs(job_queue, cancelled_ids_set, max_batch_size, timeout=1):
    """Return a list of jobs (or empty list) that can be batched together"""
    batchable_jobs = []
    required_job_info = None
    while len(batchable_jobs) < max_batch_size:
        try:
            job = job_queue.get(timeout=timeout)
        except queue.Empty:
            break
        # If pipe breaks, also gracefully return
        except OSError:
            break
        if job.job_id in cancelled_ids_set:
            # This job was cancelled, so discard it completely
            continue
        job_info = get_job_batchable_information(job)
        if required_job_info is None:
            required_job_info = job_info
        elif required_job_info != job_info:
            # This job is not batchable, put it back
            job_queue.put(job)
            # we assume here that, generally, the others jobs would also be
            # invalid. So we just return the batchable jobs we have instead
            # of going through more.
            break
        batchable_jobs.append(job)
    return batchable_jobs


def update_cancelled_jobs(cancelled_ids_queue, cancelled_ids_deque, cancelled_ids_set):
    """IN-PLACE Update cancelled_ids_set with new ids from the queue"""
    has_changed = False
    while not cancelled_ids_queue.empty():
        try:
            cancelled_id = cancelled_ids_queue.get_nowait()
        except queue.Empty:
            break
        cancelled_ids_deque.append(cancelled_id)
        has_changed = True

    if has_changed:
        cancelled_ids_set.clear()
        cancelled_ids_set.update(cancelled_ids_deque)


def predict_step(context_data, sampling_settings, model, tokens=None):
    with th.no_grad():
        predicted_step = model.predict_next_step(context_data, min_tokens_to_keep=1, tokens=tokens, **sampling_settings)
    return predicted_step


def dreamer_worker(job_queue, result_queue, cancelled_jobs_queue, quit_flag, device_to_use, args):
    logger = logging.getLogger(f"dreamer_worker {device_to_use}")
    logger.info("Loading up model...")
    model = setup_and_load_model_be_model(args)
    model = model.to(device_to_use)
    logger.info("Model loaded. Fetching results")

    cancelled_ids_deque = deque(maxlen=MAX_CANCELLED_ID_QUEUE_SIZE)
    cancelled_ids_set = set()

    while not quit_flag.is_set():
        update_cancelled_jobs(cancelled_jobs_queue, cancelled_ids_deque, cancelled_ids_set)
        batchable_jobs = fetch_list_of_batchable_jobs(job_queue, cancelled_ids_set, max_batch_size=args.max_batch_size)
        if len(batchable_jobs) == 0:
            continue
        sampling_settings = batchable_jobs[0].sampling_settings
        # make better way for passing these arguments around. sampling_settings
        # is passed as kwargs to predicting step, but max_context_length is not part of valid
        # keys there, so we need to pop it out.
        max_context_length = sampling_settings.pop("max_context_length")

        images = [job.context_images[:, :max_context_length] for job in batchable_jobs]
        actions = [job.context_actions[:, :max_context_length] for job in batchable_jobs]
        tokens = [job.context_tokens for job in batchable_jobs]

        images = th.concat(images, dim=0).to(device_to_use)
        actions = th.concat(actions, dim=0).to(device_to_use)

        context_data = TensorDict({
            "images": images,
            "actions_output": actions
        }, batch_size=images.shape[:2])

        predicted_step, predicted_image_tokens = predict_step(context_data, sampling_settings, model, tokens)

        predicted_step = predicted_step.cpu()
        predicted_images = predicted_step["images"]
        predicted_actions = predicted_step["actions_output"]
        predicted_image_tokens = predicted_image_tokens.cpu()

        for job_i, job in enumerate(batchable_jobs):
            image_context = job.context_images
            action_context = job.context_actions
            token_context = job.context_tokens
            # Keep batch dimension
            dreamt_image = predicted_images[job_i].unsqueeze(0)
            dreamt_action = predicted_actions[job_i].unsqueeze(0)
            dreamt_tokens = predicted_image_tokens[job_i].unsqueeze(0)

            # Replace the dreamed action if provided
            actions_to_take = job.actions_to_take
            if actions_to_take is not None and actions_to_take.shape[1] > 0:
                dreamt_action = actions_to_take[:, 0:1]
                # Remove the action we took
                actions_to_take = actions_to_take[:, 1:]
                if actions_to_take.shape[1] == 0:
                    actions_to_take = None

            result_queue.put(DreamJobResult(
                job_id=job.job_id,
                dream_step_index=job.num_predictions_done,
                dreamt_image=dreamt_image,
                dreamt_action=dreamt_action,
                dreamt_tokens=dreamt_tokens
            ))

            # Add job back in the queue if we have more steps to do
            if job.num_predictions_remaining > 0:
                # Stack the dreamt image and action to the context
                if image_context.shape[1] >= max_context_length:
                    image_context = image_context[:, 1:]
                    action_context = action_context[:, 1:]
                    token_context = token_context[1:]
                image_context = th.cat([image_context, dreamt_image], dim=1)
                action_context = th.cat([action_context, dreamt_action], dim=1)
                token_context.append(dreamt_tokens[0, 0].tolist())
                # We need to add context length back to sampling settings...
                # add some better way of passing these settings around
                job.sampling_settings["max_context_length"] = max_context_length
                job_queue.put(DreamJob(
                    job_id=job.job_id,
                    sampling_settings=job.sampling_settings,
                    num_predictions_remaining=job.num_predictions_remaining - 1,
                    num_predictions_done=job.num_predictions_done + 1,
                    context_images=image_context,
                    context_actions=action_context,
                    context_tokens=token_context,
                    actions_to_take=actions_to_take
                ))


class DreamerServer:
    def __init__(self, num_workers, args):
        self.num_workers = num_workers
        self.args = args
        self.model = None
        self.jobs = mp.Queue(maxsize=args.max_concurrent_jobs)
        self.results_queue = mp.Queue()
        self.cancelled_jobs = set()
        self.cancelled_jobs_queues = [mp.Queue() for _ in range(num_workers)]
        # job_id -> results
        self._last_result_cleanup = datetime.now()
        self._max_job_lifespan_datetime = timedelta(seconds=args.max_job_lifespan)
        self.local_results = defaultdict(list)
        self.logger = logging.getLogger("DreamerServer")

    def get_details(self):
        details = {
            "model_file": self.args.model,
            "max_concurrent_jobs": self.args.max_concurrent_jobs,
            "max_dream_steps_per_job": self.args.max_dream_steps_per_job,
            "max_job_lifespan": self.args.max_job_lifespan,
        }
        return json.dumps(details)

    def _check_if_should_remove_old_jobs(self):
        time_now = datetime.now()
        # Only cleanup every JOB_CLEANUP_CHECK_RATE seconds at most
        if time_now - self._last_result_cleanup < JOB_CLEANUP_CHECK_RATE:
            return

        self._last_result_cleanup = time_now
        # First add existing results to the local results
        self._gather_new_results()
        # Check if we should remove old jobs
        job_ids = list(self.local_results.keys())
        for job_id in job_ids:
            results = self.local_results[job_id]
            # If newest result is older than max_job_lifespan, remove the job
            if time_now - results[-1].result_creation_time > self._max_job_lifespan_datetime:
                self.logger.info(f"Deleted job {job_id} because it was too old. Last result was {results[-1].result_creation_time}")
                del self.local_results[job_id]

    def add_new_job(self, request, request_json):
        """
        Add new dreaming job to the queues.
        Request should have:


        Returns: json object with new job id
        """
        self._check_if_should_remove_old_jobs()

        sampling_settings = copy.deepcopy(DEFAULT_SAMPLING_SETTINGS)
        if "num_steps_to_predict" not in request_json:
            return make_response("num_steps_to_predict not in request", 400)
        num_steps_to_predict = request_json['num_steps_to_predict']
        if num_steps_to_predict > self.args.max_dream_steps_per_job:
            return make_response(f"num_steps_to_predict too large. Max {self.args.max_dream_steps_per_job}", 400)

        num_parallel_predictions = int(request_json['num_parallel_predictions']) if 'num_parallel_predictions' in request_json else 1

        if (self.jobs.qsize() + num_parallel_predictions) >= self.args.max_concurrent_jobs:
            return make_response(f"Too many jobs already running. Max {self.args.max_concurrent_jobs}", 400)

        for key in sampling_settings:
            sampling_settings[key] = float_or_none(request_json[key]) if key in request_json else sampling_settings[key]

        context_images = []
        context_actions = []
        context_tokens = []
        future_actions = []

        for step in request_json["steps"]:
            image_path = step["image_name"]
            image = np.array(Image.open(request.files[image_path].stream))
            image = be_image_preprocess(image, target_width=self.args.image_width, target_height=self.args.image_height)
            context_images.append(th.from_numpy(image))

            action = step["action"]
            action = action_vector_to_be_action_vector(action)
            context_actions.append(th.tensor(action))

            tokens = step["tokens"]
            context_tokens.append(tokens)

        future_actions = None
        if "future_actions" in request_json:
            future_actions = []
            for step in request_json["future_actions"]:
                # The rest is the action vector
                action = step["action"]
                action = action_vector_to_be_action_vector(action)
                # Add sequence and batch dimensions
                future_actions.append(th.tensor(action))

        # Add batch dimensions
        context_images = th.stack(context_images).unsqueeze(0)
        context_actions = th.stack(context_actions).unsqueeze(0)
        future_actions = th.stack(future_actions).unsqueeze(0) if future_actions is not None else None

        list_of_job_ids = []
        for _ in range(num_parallel_predictions):
            job_id = uuid.uuid4().hex
            self.jobs.put(DreamJob(
                job_id=job_id,
                sampling_settings=sampling_settings,
                num_predictions_remaining=num_steps_to_predict,
                num_predictions_done=0,
                context_images=context_images,
                context_actions=context_actions,
                context_tokens=context_tokens,
                actions_to_take=future_actions
            ))
            list_of_job_ids.append(job_id)

        job_queue_size = self.jobs.qsize()
        return json.dumps({"job_ids": list_of_job_ids, "current_jobs_in_queue": job_queue_size})

    def _gather_new_results(self):
        if not self.results_queue.empty():
            for _ in range(self.results_queue.qsize()):
                result = self.results_queue.get()
                if result.job_id in self.cancelled_jobs:
                    # Discard result if job was cancelled
                    continue
                self.local_results[result.job_id].append(result)

    def get_new_results(self, request, request_json):
        if "job_ids" not in request_json:
            return make_response("job_ids not in request", 400)
        self._gather_new_results()
        job_ids = request_json["job_ids"]
        if not isinstance(job_ids, list):
            job_ids = [job_ids]
        return_results = []
        for job_id in job_ids:
            if job_id in self.local_results:
                return_results.append(self.local_results[job_id])
                del self.local_results[job_id]

        if len(return_results) == 0:
            return make_response("No new responses", 204)

        output_json = []
        output_image_bytes = {}
        for job_results in return_results:
            for result in job_results:
                action = result.dreamt_action.numpy()
                # Remember to remove batch and sequence dimensions
                action = be_action_vector_to_action_vector(action[0, 0].tolist())
                dreamt_tokens = result.dreamt_tokens[0, 0].tolist()
                image_filename = f"{result.job_id}_{result.dream_step_index}.png"
                output_json.append({
                    "job_id": result.job_id,
                    "dream_step_index": result.dream_step_index,
                    "action": action,
                    "tokens": dreamt_tokens,
                    "image_filename": image_filename
                })

                image_bytes = io.BytesIO()
                # this probably is not as smooth as it could be
                T.ToPILImage()(result.dreamt_image[0, 0]).save(image_bytes, format="PNG")
                output_image_bytes[image_filename] = image_bytes.getvalue()

        # Write a zip file with all the images
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
        zip_bytes = io.BytesIO()
        with zipfile.ZipFile(zip_bytes, "w") as z:
            for filename, bytes in output_image_bytes.items():
                z.writestr(filename, bytes)
            # Write the json
            z.writestr(PREDICTION_JSON_FILENAME, json.dumps(output_json))

        zip_bytes.seek(0)

        return send_file(
            zip_bytes,
            mimetype="zip",
            as_attachment=True,
            download_name=f"dreaming_results_{timestamp}.zip"
        )

    def cancel_job(self, request, request_json):
        if "job_id" not in request_json:
            return make_response("job_id not in request", 400)
        job_id = request_json["job_id"]
        self.cancelled_jobs.add(job_id)
        # Cancel all jobs in the queue with this id
        for job_queue in self.cancelled_jobs_queues:
            job_queue.put(job_id)
        return make_response("OK", 200)


def main_run(args):
    app = Flask(__name__)

    num_workers = th.cuda.device_count()
    if num_workers == 0:
        raise RuntimeError("No CUDA devices found. Cannot run Dreamer.")

    server = DreamerServer(num_workers, args)
    quit_flag = mp.Event()

    # Start the dreamer worker(s)
    dreamer_worker_processes = []
    for device_i in range(num_workers):
        device = f"cuda:{device_i}"
        dreamer_worker_process = mp.Process(
            target=dreamer_worker,
            args=(server.jobs, server.results_queue, server.cancelled_jobs_queues[device_i], quit_flag, device, args)
        )
        dreamer_worker_process.daemon = True
        dreamer_worker_process.start()
        dreamer_worker_processes.append(dreamer_worker_process)

    # Add the API endpoints
    @app.route('/')
    def details():
        return server.get_details()

    @app.route('/new_job', methods=['POST'])
    def new_job():
        request_json = json.loads(request.form["json"])
        return server.add_new_job(request, request_json)

    @app.route('/get_job_results', methods=['GET'])
    def get_results():
        # the "Json" is now in regular GET payload/parameters
        request_json = {"job_ids": request.args.getlist("job_ids")}
        return server.get_new_results(request, request_json)

    @app.route('/cancel_job', methods=['GET'])
    def cancel_job():
        request_json = request.args.to_dict()
        return server.cancel_job(request, request_json)

    app.run(host="0.0.0.0", port=args.port, debug=args.debug)

    # Cleanup
    quit_flag.set()
    for dreamer_worker_process in dreamer_worker_processes:
        dreamer_worker_process.join()


if __name__ == '__main__':
    args = parser.parse_args()
    main_run(args)