Spaces:
Runtime error
Runtime error
| import collections | |
| import random | |
| import time | |
| import multiprocessing as mp | |
| import json | |
| from PIL import Image | |
| from compiled_jss.CPEnv import CompiledJssEnvCP | |
| from stable_baselines3.common.vec_env import VecEnvWrapper | |
| from torch.distributions import Categorical | |
| import torch | |
| import numpy as np | |
| from MyVecEnv import WrapperRay | |
| import gradio as gr | |
| import docplex.cp.utils_visu as visu | |
| import matplotlib.pyplot as plt | |
| class VecPyTorch(VecEnvWrapper): | |
| def __init__(self, venv, device): | |
| super(VecPyTorch, self).__init__(venv) | |
| self.device = device | |
| def reset(self): | |
| return self.venv.reset() | |
| def step_async(self, actions): | |
| self.venv.step_async(actions) | |
| def step_wait(self): | |
| return self.venv.step_wait() | |
| def make_env(seed, instance): | |
| def thunk(): | |
| _env = CompiledJssEnvCP(instance) | |
| return _env | |
| return thunk | |
| def solve(file): | |
| random.seed(0) | |
| np.random.seed(0) | |
| torch.manual_seed(0) | |
| num_workers = min(mp.cpu_count(), 32) | |
| with torch.inference_mode(): | |
| device = torch.device('cpu') | |
| actor = torch.jit.load('actor.pt', map_location=device) | |
| actor.eval() | |
| start_time = time.time() | |
| fn_env = [make_env(0, file.name) | |
| for _ in range(num_workers)] | |
| ray_wrapper_env = WrapperRay(lambda n: fn_env[n](), | |
| num_workers, 1, device) | |
| envs = VecPyTorch(ray_wrapper_env, device) | |
| current_solution_cost = float('inf') | |
| current_solution = '' | |
| obs = envs.reset() | |
| total_episode = 0 | |
| while total_episode < envs.num_envs: | |
| logits = actor(obs['interval_rep'], obs['attention_interval_mask'], obs['job_resource_mask'], | |
| obs['action_mask'], obs['index_interval'], obs['start_end_tokens']) | |
| # temperature vector | |
| if num_workers >= 4: | |
| temperature = torch.arange(0.5, 2.0, step=(1.5 / num_workers), device=device) | |
| else: | |
| temperature = torch.ones(num_workers, device=device) | |
| logits = logits / temperature[:, None] | |
| probs = Categorical(logits=logits).probs | |
| # random sample based on logits | |
| actions = torch.multinomial(probs, probs.shape[1]).cpu().numpy() | |
| obs, reward, done, infos = envs.step(actions) | |
| total_episode += done.sum() | |
| # total_actions += 1 | |
| # print(f'Episode {total_episode} / {envs.num_envs} - Actions {total_actions}', end='\r') | |
| for env_idx, info in enumerate(infos): | |
| if 'makespan' in info and int(info['makespan']) < current_solution_cost: | |
| current_solution_cost = int(info['makespan']) | |
| current_solution = json.loads(info['solution']) | |
| total_time = time.time() - start_time | |
| pretty_output = "" | |
| for job_id in range(len(current_solution)): | |
| pretty_output += f"Job {job_id}: {current_solution[job_id]}\n" | |
| jobs_data = [] | |
| file.seek(0) | |
| line_str: str = file.readline() | |
| line_cnt: int = 1 | |
| while line_str: | |
| data = [] | |
| split_data = line_str.split() | |
| if line_cnt == 1: | |
| jobs_count, machines_count = int(split_data[0]), int( | |
| split_data[1] | |
| ) | |
| else: | |
| i = 0 | |
| this_job_op_count = 0 | |
| while i < len(split_data): | |
| machine, op_time = int(split_data[i]), int(split_data[i + 1]) | |
| data.append((machine, op_time)) | |
| i += 2 | |
| this_job_op_count += 1 | |
| jobs_data.append(data) | |
| line_str = file.readline() | |
| line_cnt += 1 | |
| visu.timeline(f'Solution for job-shop, solved using ') | |
| visu.panel('Jobs') | |
| # convert to integer the current_solution | |
| current_solution = [[int(x) for x in y] for y in current_solution] | |
| for job_id in range(len(current_solution)): | |
| visu.sequence(name=f'J{job_id}', intervals=[(current_solution[job_id][task_id], | |
| current_solution[job_id][task_id] + jobs_data[job_id][task_id][ | |
| 1], jobs_data[job_id][task_id][0], | |
| f'M{jobs_data[job_id][task_id][0]}') | |
| for task_id in | |
| range(len(current_solution[job_id]))]) | |
| visu.panel('Machines') | |
| machine_solution = collections.defaultdict(list) | |
| for job_id in range(len(current_solution)): | |
| for task_id in range(len(current_solution[job_id])): | |
| machine = jobs_data[job_id][task_id][1] | |
| machine_solution[machine].append((current_solution[job_id][task_id], | |
| current_solution[job_id][task_id] + jobs_data[job_id][task_id][1], | |
| machine, f'J{job_id}')) | |
| # sort dictionary keys | |
| machine_solution = {k: machine_solution[k] for k in sorted(machine_solution.keys())} | |
| for machine_id in machine_solution: | |
| visu.sequence(name=f'M{machine_id}', | |
| intervals=machine_solution[machine_id]) | |
| plt.rcParams["font.family"] = "Times New Roman" | |
| plt.rcParams["font.size"] = "30" | |
| plt.gca().set_aspect('equal') | |
| plt.rcParams["figure.figsize"] = (45, 50) | |
| from io import BytesIO | |
| buffer = BytesIO() | |
| visu.show(pngfile=buffer) | |
| reloadedPILImage = Image.open(buffer) | |
| return pretty_output, reloadedPILImage, str(total_time) + " seconds" | |
| title = "Job-Shop Scheduling CP RL" | |
| description = "A Job-Shop Scheduling Reinforcement Learning based solver, using an underlying CP model as an " \ | |
| "environment. " | |
| article = "<p style='text-align: center'>Article Under Review</p>" | |
| examples = ['ta01', 'dmu01.txt', 'la01.txt'] | |
| iface = gr.Interface(fn=solve, inputs=gr.File(label="Instance File"), outputs=[gr.Text(label="Solution"), gr.Image(label="Solution's Gantt Chart"), gr.Text(label="Elapsed Time")], title=title, description=description, article=article, examples=examples) | |
| iface.launch(enable_queue=True) | |