| import os |
| from enum import Enum |
|
|
| import typer |
| from openai import OpenAI |
|
|
| from jssp_openenv.client import JSSPEnvClient |
| from jssp_openenv.gantt import gantt_chart |
| from jssp_openenv.policy import JSSPEnvPolicy, JSSPFifoPolicy, JSSPLLMPolicy, JSSPMaxMinPolicy |
| from jssp_openenv.solver import solve_jssp |
|
|
| SERVER_URL = "http://localhost:8000" |
| MAX_STEPS = 1000 |
| OUTPUT_DIR = "output" |
| os.makedirs(OUTPUT_DIR, exist_ok=True) |
|
|
| cli = typer.Typer() |
|
|
|
|
| class PolicyName(str, Enum): |
| FIFO = "fifo" |
| LLM = "llm" |
| MAX_MIN = "maxmin" |
|
|
|
|
| @cli.command() |
| def solve( |
| policy: PolicyName = typer.Argument(help="The policy to use"), |
| server_url: str = typer.Option(SERVER_URL, help="The URL of the JSSP server"), |
| max_steps: int = typer.Option(MAX_STEPS, help="The maximum number of steps per instance"), |
| verbose: bool = typer.Option(False, "--verbose", "-v", help="Whether to print verbose output"), |
| model_id: str = typer.Option(None, "--model-id", "-m", help="The ID of the model to use"), |
| ): |
| """Solve a JSSP instance using the given policy.""" |
| env_client = JSSPEnvClient(base_url=server_url) |
|
|
| policy_obj: JSSPEnvPolicy |
| match policy: |
| case PolicyName.FIFO: |
| policy_obj = JSSPFifoPolicy() |
| title = "FIFO Policy" |
| filename = "gantt_fifo_policy.png" |
|
|
| case PolicyName.LLM: |
| if not model_id: |
| raise ValueError("You must set --model-id to use the LLM policy") |
| api_key = os.getenv("HF_TOKEN") |
| if not api_key: |
| raise ValueError("You must set the HF_TOKEN environment variable to use the LLM policy") |
| client = OpenAI(base_url="https://router.huggingface.co/v1", api_key=api_key) |
| policy_obj = JSSPLLMPolicy(client=client, model_id=model_id) |
| title = f"LLM Policy ({model_id})" |
| filename = f"gantt_llm_policy_{model_id.replace('/', '_').replace(':', '_').replace('-', '_').replace(' ', '_')}.png" |
|
|
| case PolicyName.MAX_MIN: |
| policy_obj = JSSPMaxMinPolicy() |
| title = "Max-Min Policy" |
| filename = "gantt_maxmin_policy.png" |
|
|
| makespan, scheduled_events = solve_jssp(env_client, policy_obj, max_steps, verbose) |
|
|
| if verbose: |
| print("Schedule events:") |
| for event in scheduled_events: |
| print( |
| f"[{event.start_time}] Scheduling job {event.job_id} on machine {event.machine_id} for {event.end_time - event.start_time} minute(s)" |
| ) |
|
|
| print(f"Solved in {makespan} steps") |
|
|
| filepath = os.path.join(OUTPUT_DIR, filename) |
| gantt_chart(scheduled_events, title=title, makespan=makespan, save_to=filepath) |
| print(f"Saved Gantt chart to {filepath}") |
|
|
|
|
| if __name__ == "__main__": |
| cli() |
|
|
