initial commit

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,219 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#   and can be added to the global gitignore or merged into this file.  For a more nuclear
+#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
+# .idea/
+
+# Abstra
+#   Abstra is an AI-powered process automation framework.
+#   Ignore directories containing user credentials, local state, and settings.
+#   Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#   you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml
+
+charts/
+!charts/gantt_fifo_policy.png

Dockerfile

@@ -0,0 +1,18 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+# Copy project files
+COPY pyproject.toml ./
+COPY src/ ./src/
+COPY app.py ./
+
+# Install the project
+RUN pip install --no-cache-dir -e .
+
+# Expose port 7860
+EXPOSE 7860
+
+# Run uvicorn directly
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+

README.md

@@ -0,0 +1,168 @@
+---
+title: JSSP OpenEnv
+emoji: ⏰
+colorFrom: green
+colorTo: purple
+sdk: docker
+pinned: false
+---
+
+<p align="center">
+  <img src="assets/jssp_openenv.png" alt="jssp_openenv" width="400">
+</p>
+
+<p align="center">
+  <a href="https://huggingface.co/spaces/Wauplin/jssp_openenv" style="font-size: 1.2em;">Try it live on Hugging Face Spaces</a>
+</p>
+
+## Job shop scheduling problem (JSSP)
+
+The [Job Shop Scheduling Problem](https://en.wikipedia.org/wiki/Job-shop_scheduling) (JSSP) is a classic optimization problem in operations research. Given a set of jobs, each consisting of multiple operations that must be performed in a specific sequence, and a set of machines, the goal is to schedule the operations on machines to minimize the total completion time (makespan).
+
+**Key constraints:**
+- Each job consists of a sequence of operations that must be completed in order
+- Each operation requires a specific machine for a given duration
+- Each machine can process only one operation at a time
+- Once started, an operation cannot be interrupted
+
+This implementation uses the OpenEnv framework to create a reinforcement learning environment where an agent (policy) learns to make scheduling decisions at each time step.
+
+> !TIP
+> For now, we only implement and run the FT06 problem. It is a well-known problem in the literature with a known optimal solution.
+> Goal for training is to run arbitrary random environments.
+
+## OpenEnv
+
+[OpenEnv](https://github.com/meta-pytorch/OpenEnv) is a framework from Meta PyTorch and Hugging Face for building reinforcement learning environments. It provides:
+
+- A standardized interface for environments with `Action` and `Observation` models
+- A web-based interface for interactive exploration of environments
+- A client-server architecture for distributed training and evaluation
+- Integration with LLM-based policies for solving complex problems
+
+This project implements a JSSP environment using OpenEnv, allowing you to:
+- Interact with the environment through a web interface
+- Test different scheduling policies (FIFO, Max-Min, LLM-based)
+- Train reinforcement learning agents to solve JSSP instances
+
+## Project Architecture
+
+The project follows a client-server architecture using the OpenEnv framework:
+
+### Core Components
+
+**Models** (`src/jssp_openenv/models.py`):
+- `JSSPAction`: Represents scheduling actions (list of job IDs to schedule)
+- `JSSPObservation`: Contains the current state (machines, ready operations, progress)
+
+**Environment** (`src/jssp_openenv/server/jssp_environment.py`):
+- `JSSPEnvironment`: The core simulation environment that:
+  - Manages job progress and machine states
+  - Validates actions and enforces constraints
+  - Advances simulation time using SimPy
+  - Returns observations and rewards
+
+**Client** (`src/jssp_openenv/client.py`):
+- `JSSPEnvClient`: HTTP client that communicates with the environment server
+- Handles action serialization and observation parsing
+
+**Policies** (`src/jssp_openenv/policy.py`):
+- `JSSPEnvPolicy`: Abstract base class for scheduling policies
+- `JSSPFifoPolicy`: First-In-First-Out scheduling (schedules jobs by ID order)
+- `JSSPMaxMinPolicy`: Max-Min scheduling (prioritizes longest operations)
+- `JSSPLLMPolicy`: LLM-based scheduling using OpenAI-compatible APIs
+
+**Solver** (`src/jssp_openenv/solver.py`):
+- `solve_jssp()`: Orchestrates the solving process by:
+  - Resetting the environment
+  - Iteratively applying policy actions
+  - Tracking scheduled events for visualization
+  - Returning makespan and event history
+
+**Visualization** (`src/jssp_openenv/gantt.py`):
+- Generates Gantt charts showing the schedule timeline
+
+## How to use
+
+### Install
+
+Install the package and its dependencies:
+
+```bash
+pip install -e .
+```
+
+For development with additional tools (pytest, ruff, etc.):
+
+```bash
+pip install -e ".[dev]"
+```
+
+**Note:** For LLM-based policies, you'll need to set the `HF_TOKEN` environment variable with your Hugging Face API token:
+
+```bash
+export HF_TOKEN=your_token_here
+```
+
+### Run server
+
+To play with the environment locally, run
+
+```
+python app.py
+```
+
+and go to http://0.0.0.0:8000/web.
+
+### Run policy
+
+**FIFO policy** (always run first available job):
+
+```
+python run.py fifo
+```
+
+**Max-Min policy** (always run longest job first):
+
+```
+python run.py maxmin
+```
+
+**LLM policy** (ask an LLM to solve the problem)
+
+```
+python run.py llm --model-id "openai/gpt-oss-20b:groq"
+python run.py llm --model-id "openai/gpt-oss-120b:cerebras"
+python run.py llm --model-id "Qwen/Qwen3-32B:groq"
+```
+
+### Check results
+
+The solver will resolve the problem using the policy and then plot a gantt chart of the solution in the `./charts` folder.
+
+Here is an example:
+
+![FIFO Policy Gantt Chart](assets/gantt_fifo_policy.png)
+
+## Run with docker
+
+Build the Docker image:
+
+```bash
+docker build -t jssp-openenv .
+```
+
+Run the container:
+
+```bash
+docker run -p 7860:7860 jssp-openenv
+```
+
+The web interface will be available at http://localhost:7860/web.
+
+## TODO
+
+- [ ] run on other example environments (FT10, FT20)
+- [ ] run on random environments
+- [ ] run multiple policies and summarize results
+- [ ] trainer

app.py

@@ -0,0 +1,13 @@
+from openenv_core.env_server import create_web_interface_app
+
+from jssp_openenv.examples import FT06
+from jssp_openenv.models import JSSPAction, JSSPObservation
+from jssp_openenv.server.jssp_environment import JSSPEnvironment
+
+env = JSSPEnvironment(FT06)
+app = create_web_interface_app(env, JSSPAction, JSSPObservation, "JSSP (FT06)")
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

pyproject.toml

@@ -0,0 +1,40 @@
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "jssp_openenv"
+description = "JSSP environment for OpenEnv"
+version = "0.1.0"
+authors = [
+    { name="Lucain Pouget", email="lucain@huggingface.co" }
+]
+requires-python = ">=3.10"
+dependencies = [
+    "openenv-core>=0.1.0",
+    "openai", # for inference.py
+    "simpy", # for env simulation
+    "matplotlib", # for plotting
+    "typer", # for CLI
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "ruff",
+    "ty",
+]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.ruff]
+exclude = [".git", ".ruff_cache", ".venv"]
+line-length = 119
+# Ignored rules:
+# "E501" -> line length violation
+lint.ignore = ["E501"]
+lint.select = ["E", "F", "I", "W"]

run.py

@@ -0,0 +1,77 @@
+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  # Maximum number of steps per instance
+CHART_DIR = "charts"
+os.makedirs(CHART_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(CHART_DIR, filename)
+    gantt_chart(scheduled_events, title=title, makespan=makespan, save_to=filepath)
+    print(f"Saved Gantt chart to {filepath}")
+
+
+if __name__ == "__main__":
+    cli()

src/jssp_openenv/client.py

@@ -0,0 +1,25 @@
+from openenv_core import HTTPEnvClient, StepResult
+
+from .models import JSSPAction, JSSPObservation, MachineObservation, ReadyOperationObservation
+
+
+class JSSPEnvClient(HTTPEnvClient[JSSPAction, JSSPObservation]):
+    def _step_payload(self, action: JSSPAction) -> dict:
+        return {"job_ids": action.job_ids}
+
+    def _parse_result(self, payload: dict) -> StepResult[JSSPObservation]:
+        obs_data = payload["observation"]
+        return StepResult[JSSPObservation](
+            observation=JSSPObservation(
+                machines=[MachineObservation(**machine) for machine in obs_data.pop("machines")],
+                ready_operations=[
+                    ReadyOperationObservation(**operation) for operation in obs_data.pop("ready_operations")
+                ],
+                **obs_data,
+            ),
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: dict) -> dict:
+        return payload

src/jssp_openenv/examples.py

@@ -0,0 +1,61 @@
+"""
+Example instances for the JSSP environment.
+
+Checkout https://github.com/tamy0612/JSPLIB for more instances.
+"""
+
+from .models import JobT
+
+# Fisher and Thompson 6x6 instance
+# https://github.com/tamy0612/JSPLIB/blob/master/instances/ft06
+# optimal solution: 55
+FT06: list[JobT] = [
+    [(2, 1), (0, 3), (1, 6), (3, 7), (5, 3), (4, 6)],
+    [(1, 8), (2, 5), (4, 10), (5, 10), (0, 10), (3, 4)],
+    [(2, 5), (3, 4), (5, 8), (0, 9), (1, 1), (4, 7)],
+    [(1, 5), (0, 5), (2, 5), (3, 3), (4, 8), (5, 9)],
+    [(2, 9), (1, 3), (4, 5), (5, 4), (0, 3), (3, 1)],
+    [(1, 3), (3, 3), (5, 9), (0, 10), (4, 4), (2, 1)],
+]
+
+# Fisher and Thompson 10x10 instance
+# https://github.com/tamy0612/JSPLIB/blob/master/instances/ft10
+# optimal solution: 930
+FT_10: list[JobT] = [
+    [(0, 29), (1, 78), (2, 9), (3, 36), (4, 49), (5, 11), (6, 62), (7, 56), (8, 44), (9, 21)],
+    [(0, 43), (2, 90), (4, 75), (9, 11), (3, 69), (1, 28), (6, 46), (5, 46), (7, 72), (8, 30)],
+    [(1, 91), (0, 85), (3, 39), (2, 74), (8, 90), (5, 10), (7, 12), (6, 89), (9, 45), (4, 33)],
+    [(1, 81), (2, 95), (0, 71), (4, 99), (6, 9), (8, 52), (7, 85), (3, 98), (9, 22), (5, 43)],
+    [(2, 14), (0, 6), (1, 22), (5, 61), (3, 26), (4, 69), (8, 21), (7, 49), (9, 72), (6, 53)],
+    [(2, 84), (1, 2), (5, 52), (3, 95), (8, 48), (9, 72), (0, 47), (6, 65), (4, 6), (7, 25)],
+    [(1, 46), (0, 37), (3, 61), (2, 13), (6, 32), (5, 21), (9, 32), (8, 89), (7, 30), (4, 55)],
+    [(2, 31), (0, 86), (1, 46), (5, 74), (4, 32), (6, 88), (8, 19), (9, 48), (7, 36), (3, 79)],
+    [(0, 76), (1, 69), (3, 76), (5, 51), (2, 85), (9, 11), (6, 40), (7, 89), (4, 26), (8, 74)],
+    [(1, 85), (0, 13), (2, 61), (6, 7), (8, 64), (9, 76), (5, 47), (3, 52), (4, 90), (7, 45)],
+]
+
+# Fisher and Thompson 20x5 instance
+# https://github.com/tamy0612/JSPLIB/blob/master/instances/ft20
+# optimal solution: 1165
+FT20: list[JobT] = [
+    [(0, 29), (1, 9), (2, 49), (3, 62), (4, 44)],
+    [(0, 43), (1, 75), (3, 69), (2, 46), (4, 72)],
+    [(1, 91), (0, 39), (2, 90), (4, 12), (3, 45)],
+    [(1, 81), (0, 71), (4, 9), (2, 85), (3, 22)],
+    [(2, 14), (1, 22), (0, 26), (3, 21), (4, 72)],
+    [(2, 84), (1, 52), (4, 48), (0, 47), (3, 6)],
+    [(1, 46), (0, 61), (2, 32), (3, 32), (4, 30)],
+    [(2, 31), (1, 46), (0, 32), (3, 19), (4, 36)],
+    [(0, 76), (3, 76), (2, 85), (1, 40), (4, 26)],
+    [(1, 85), (2, 61), (0, 64), (3, 47), (4, 90)],
+    [(1, 78), (3, 36), (0, 11), (4, 56), (2, 21)],
+    [(2, 90), (0, 11), (1, 28), (3, 46), (4, 30)],
+    [(0, 85), (2, 74), (1, 10), (3, 89), (4, 33)],
+    [(2, 95), (0, 99), (1, 52), (3, 98), (4, 43)],
+    [(0, 6), (1, 61), (4, 69), (2, 49), (3, 53)],
+    [(1, 2), (0, 95), (3, 72), (4, 65), (2, 25)],
+    [(0, 37), (2, 13), (1, 21), (3, 89), (4, 55)],
+    [(0, 86), (1, 74), (4, 88), (2, 48), (3, 79)],
+    [(1, 69), (2, 51), (0, 11), (3, 89), (4, 74)],
+    [(0, 13), (1, 7), (2, 76), (3, 52), (4, 45)],
+]

src/jssp_openenv/gantt.py

@@ -0,0 +1,77 @@
+import matplotlib.pyplot as plt
+
+from .models import ScheduledEvent
+
+
+def gantt_chart(scheduled_events: list[ScheduledEvent], title: str, makespan: int, save_to: str) -> None:
+    """Generate and save a Gantt chart from schedule events using matplotlib."""
+    if not scheduled_events:
+        print("No schedule events to save.")
+        return
+
+    # Extract unique machines and jobs
+    machines = sorted(set(event.machine_id for event in scheduled_events))
+    jobs = sorted(set(event.job_id for event in scheduled_events))
+
+    # Create figure and axis
+    fig, ax = plt.subplots(figsize=(12, max(6, len(machines) * 0.8)))
+
+    # Color map for different jobs
+    colors = plt.cm.tab20(range(len(jobs)))
+    job_color_map = {job_id: colors[i % len(colors)] for i, job_id in enumerate(jobs)}
+
+    # Track which jobs have been added to legend
+    legend_added = set()
+
+    # Plot each schedule event as a horizontal bar
+    for event in scheduled_events:
+        machine_idx = machines.index(event.machine_id)
+        duration = event.end_time - event.start_time
+
+        # Only add label for legend if this job hasn't been added yet
+        label = f"Job {event.job_id}" if event.job_id not in legend_added else ""
+        if label:
+            legend_added.add(event.job_id)
+
+        ax.barh(
+            machine_idx,
+            duration,
+            left=event.start_time,
+            height=0.6,
+            color=job_color_map[event.job_id],
+            edgecolor="black",
+            linewidth=0.5,
+            label=label,
+        )
+
+        # Add job label in the middle of the bar
+        mid_time = event.start_time + duration / 2
+        ax.text(
+            mid_time,
+            machine_idx,
+            f"J{event.job_id}",
+            ha="center",
+            va="center",
+            fontsize=8,
+            fontweight="bold",
+            color="white" if sum(job_color_map[event.job_id][:3]) < 1.5 else "black",
+        )
+
+    # Customize the chart
+    ax.set_yticks(range(len(machines)))
+    ax.set_yticklabels([f"Machine {m}" for m in machines])
+    ax.set_xlabel("Time", fontsize=12)
+    ax.set_ylabel("Machine", fontsize=12)
+    ax.set_title(f"{title} (Makespan: {makespan})", fontsize=14, fontweight="bold")
+    ax.grid(True, axis="x", alpha=0.3, linestyle="--")
+
+    # Set x-axis limits with some padding
+    max_time = max(event.end_time for event in scheduled_events) if scheduled_events else 0
+    ax.set_xlim(0, max_time * 1.05)
+
+    # Add legend
+    ax.legend(loc="upper right", title="Jobs")
+
+    plt.tight_layout()
+    plt.savefig(save_to)
+    plt.close(fig)

src/jssp_openenv/models.py

@@ -0,0 +1,74 @@
+"""
+Data models for the JSSP Environment.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+from openenv_core import Action, Observation
+
+JobT = list[tuple[int, int]]  # (machine_index, processing_time)
+
+
+@dataclass(kw_only=True)
+class JSSPAction(Action):
+    """Action for the JSSP environment."""
+
+    job_ids: list[int]
+
+    def __post_init__(self):
+        if isinstance(self.job_ids, str):
+            # For web app
+            self.job_ids = parse_job_ids(self.job_ids)
+
+
+@dataclass(kw_only=True)
+class MachineObservation:
+    """Observation of a single machine in the JSSP environment."""
+
+    machine_id: int
+    busy_until: Optional[int]
+    current_job_id: Optional[int]
+
+
+@dataclass
+class ReadyOperationObservation:
+    job_id: int
+    machine_id: int
+    duration: int
+    remaining_ops: int
+
+
+@dataclass(kw_only=True)
+class JSSPObservation(Observation):
+    """Observation from the JSSP environment - the echoed message."""
+
+    episode_id: str
+
+    step_count: int
+    machines: list[MachineObservation]
+    ready_operations: list[ReadyOperationObservation]
+    completed_jobs: int
+    total_jobs: int
+
+
+def parse_job_ids(job_ids: str) -> list[int]:
+    """Parse job_ids from string (error out if cannot be parsed)."""
+    try:
+        return [int(job_id) for job_id in job_ids.split(",") if job_id.strip()]
+    except ValueError:
+        raise ValueError(f"Invalid job_ids: {job_ids}")
+
+
+@dataclass
+class ScheduledEvent:
+    """Represents a scheduled operation on a machine.
+
+    Used for plotting the schedule.
+    Not used for the environment / policy / solver.
+    """
+
+    job_id: int
+    machine_id: int
+    start_time: int
+    end_time: int

src/jssp_openenv/policy.py

@@ -0,0 +1,261 @@
+import re
+from abc import ABC, abstractmethod
+
+from openai import OpenAI
+
+from .models import JSSPAction, JSSPObservation, MachineObservation, ReadyOperationObservation
+
+
+class JSSPEnvPolicy(ABC):
+    @abstractmethod
+    def act(self, observation: JSSPObservation) -> JSSPAction:
+        """Act based on the observation."""
+
+
+class JSSPFifoPolicy(JSSPEnvPolicy):
+    def act(self, observation: JSSPObservation) -> JSSPAction:
+        """
+        FIFO scheduling: schedule ready operations in order of job_id.
+
+        This policy schedules operations in FIFO order (by job_id), respecting
+        machine availability. It only schedules operations for machines that are
+        currently available (not busy).
+        """
+        # Create a lookup for machine availability
+        machine_available = {m.machine_id: m.busy_until is None for m in observation.machines}
+
+        # Filter to only ready operations with available machines
+        available_ops = [op for op in observation.ready_operations if machine_available.get(op.machine_id, False)]
+
+        # Sort by job_id (FIFO: first job_id first)
+        available_ops.sort(key=lambda op: op.job_id)
+
+        # Track which machines we've already scheduled to avoid conflicts
+        scheduled_machines = set()
+        scheduled_job_ids = []
+
+        # Schedule jobs in FIFO order, but skip if machine is already taken
+        for op in available_ops:
+            if op.machine_id not in scheduled_machines:
+                scheduled_job_ids.append(op.job_id)
+                scheduled_machines.add(op.machine_id)
+
+        return JSSPAction(job_ids=scheduled_job_ids)
+
+
+class JSSPMaxMinPolicy(JSSPEnvPolicy):
+    def act(self, observation: JSSPObservation) -> JSSPAction:
+        """
+        Max-Min scheduling: schedule the operation with the longest duration first.
+        """
+        # Sort operations by duration (max-min)
+        ops = sorted(observation.ready_operations, key=lambda op: op.duration, reverse=True)
+
+        # Track which machines we've already scheduled to avoid conflicts
+        scheduled_machines = set()
+        scheduled_job_ids = []
+
+        # Schedule jobs in max-min order, but skip if machine is already taken
+        for op in ops:
+            if op.machine_id not in scheduled_machines:
+                scheduled_job_ids.append(op.job_id)
+                scheduled_machines.add(op.machine_id)
+
+        return JSSPAction(job_ids=scheduled_job_ids)
+
+
+PROMPT_TEMPLATE = """
+You are solving a Job Shop Scheduling Problem (JSSP). Your goal is to minimize the total completion time (makespan) by efficiently scheduling job operations across machines.
+
+You must optimize for minimal makespan while respecting all constraints. Each job consists of multiple operations that must be completed in sequence, and each operation requires a specific machine for a given duration.
+
+---
+
+### 🕒 Current State
+**Step:** {step_count} | **Completed:** {completed_jobs}/{total_jobs}
+
+---
+
+### ⚙️ Machine Status
+{machines_status}
+
+You must check machine availability before scheduling. Machines that are busy cannot start new operations until they finish their current task.
+
+---
+
+### ✅ Ready to Schedule (NOW)
+{ready_operations_list}
+
+Each entry shows: **machine**, **duration**, and **remaining ops**.
+You can only schedule operations that are ready at this step. These are operations whose previous steps in the job sequence have been completed.
+
+---
+
+### 🎯 Rules You Must Follow
+1. You must schedule only **ready** operations. Do not attempt to schedule operations that are not ready.
+2. Each machine can run **one job at a time**. You cannot schedule multiple jobs on the same machine simultaneously.
+3. You must not schedule jobs on **busy** machines (`busy_until > current step`). Check machine availability before scheduling.
+4. You may **schedule multiple** jobs on different machines in the same step, or you may choose to wait if no good scheduling opportunity exists.
+
+---
+
+### 🧩 Available Actions
+{legal_actions}
+
+These are the valid job IDs you can schedule at this step. You must choose from this list.
+
+**Answer format:**
+- To schedule jobs: `"0,2"` or `"1"` (comma-separated job IDs)
+- To wait: `""` (empty string)
+
+Respond only with the action format specified above.
+"""
+
+
+class JSSPLLMPolicy(JSSPEnvPolicy):
+    """LLM-based scheduling policy using OpenAI-compatible API."""
+
+    # Job Shop Scheduling prompt template
+
+    def __init__(self, client: OpenAI, model_id: str):
+        """
+        Initialize the LLM policy.
+
+        Args:
+            client: OpenAI-compatible client instance
+            model_id: Name of the model to use
+        """
+        self.client = client
+        self.model_id = model_id
+
+    def act(self, observation: JSSPObservation) -> JSSPAction:
+        """
+        LLM scheduling: use an LLM to schedule the operations.
+
+        Determines legal actions (ready operations with available machines),
+        formats a prompt, calls the LLM, and parses the response to return
+        a scheduling action.
+        """
+        # Determine machine availability
+        machine_available = {
+            m.machine_id: m.busy_until is None or m.busy_until <= observation.step_count for m in observation.machines
+        }
+
+        # Filter ready operations to only include those with available machines
+        legal_job_ids = [
+            op.job_id for op in observation.ready_operations if machine_available.get(op.machine_id, False)
+        ]
+
+        # If no legal actions, return empty action (wait)
+        if not legal_job_ids:
+            return JSSPAction(job_ids=[])
+
+        # Format prompt
+        machines_status = self._format_machines_status(observation.machines, observation.step_count)
+        ready_operations_list = self._format_ready_operations(observation.ready_operations)
+
+        prompt = PROMPT_TEMPLATE.format(
+            step_count=observation.step_count,
+            completed_jobs=observation.completed_jobs,
+            total_jobs=observation.total_jobs,
+            machines_status=machines_status,
+            ready_operations_list=ready_operations_list,
+            legal_actions=legal_job_ids,
+        )
+        print(f"Step {observation.step_count}")
+
+        # Call LLM
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model_id, messages=[{"role": "user", "content": prompt}], temperature=0.0
+            )
+            llm_output = response.choices[0].message.content or ""
+            print(f"LLM Output: {llm_output}")
+            job_ids = self._parse_action(llm_output, legal_job_ids)
+            print(f"Job IDs: {job_ids}")
+
+            # Ensure we don't schedule multiple jobs on the same machine
+            # Track which machines we've already scheduled to avoid conflicts
+            scheduled_machines = set()
+            filtered_job_ids = []
+            for job_id in job_ids:
+                # Find the operation for this job
+                op = next((op for op in observation.ready_operations if op.job_id == job_id), None)
+                if op is not None and op.machine_id not in scheduled_machines:
+                    filtered_job_ids.append(job_id)
+                    scheduled_machines.add(op.machine_id)
+
+            return JSSPAction(job_ids=filtered_job_ids)
+
+        except Exception as e:
+            print(f"Error calling LLM: {e}")
+            print(f"Prompt: {prompt}")
+            # On error, fall back to empty action (wait)
+            return JSSPAction(job_ids=[])
+
+    @staticmethod
+    def _format_machines_status(machines: list[MachineObservation], current_step: int) -> str:
+        """Format machine status for prompt."""
+        lines = []
+        for machine in machines:
+            if machine.busy_until is None or machine.busy_until <= current_step:
+                status = "Available"
+            else:
+                status = f"Busy until t={machine.busy_until}"
+            job_info = f" (job {machine.current_job_id})" if machine.current_job_id is not None else ""
+            lines.append(f"  Machine {machine.machine_id}: {status}{job_info}")
+        return "\n".join(lines) if lines else "  (No machines)"
+
+    @staticmethod
+    def _format_ready_operations(ready_operations: list[ReadyOperationObservation]) -> str:
+        """Format ready operations for prompt."""
+        lines = []
+        for op in ready_operations:
+            lines.append(
+                f"  Job {op.job_id}: Machine {op.machine_id}, Duration {op.duration} min, {op.remaining_ops} ops remaining"
+            )
+        return "\n".join(lines) if lines else "  (No ready operations)"
+
+    @staticmethod
+    def _parse_action(text: str, legal_job_ids: list[int]) -> list[int]:
+        """Parse comma-separated job IDs from model output."""
+        # First, we remove the reasoning section
+        text = text.split("<think>")[-1].split("</think>")[-1].strip()
+
+        # First, try to split by comma and extract numbers from each part
+        # This handles "2,3" or "2, 3" correctly
+        parts = text.split(",")
+        job_ids = []
+
+        # Extract numbers from each part (handles "2" or "job 2" or " 2 ")
+        for part in parts:
+            numbers = re.findall(r"\d+", part.strip())
+            for num_str in numbers:
+                try:
+                    job_id = int(num_str)
+                    if job_id in legal_job_ids:
+                        job_ids.append(job_id)
+                except ValueError:
+                    continue
+
+        # If no comma-separated values found, try extracting all numbers
+        # (handles cases like "Schedule jobs 2 and 3")
+        if not job_ids:
+            numbers = re.findall(r"\d+", text)
+            for num_str in numbers:
+                try:
+                    job_id = int(num_str)
+                    if job_id in legal_job_ids:
+                        job_ids.append(job_id)
+                except ValueError:
+                    continue
+
+        # Remove duplicates while preserving order
+        seen = set()
+        unique_job_ids = []
+        for job_id in job_ids:
+            if job_id not in seen:
+                seen.add(job_id)
+                unique_job_ids.append(job_id)
+
+        return unique_job_ids if unique_job_ids else []  # Return empty list if no valid jobs found

src/jssp_openenv/server/app.py

@@ -0,0 +1,9 @@
+from openenv_core.env_server import create_app
+
+from ..examples import FT06
+from ..models import JSSPAction, JSSPObservation
+from .jssp_environment import JSSPEnvironment
+
+# Create FastAPI app
+env = JSSPEnvironment(FT06)
+app = create_app(env, JSSPAction, JSSPObservation)

src/jssp_openenv/server/jssp_environment.py

@@ -0,0 +1,218 @@
+import uuid
+from copy import deepcopy
+from typing import Optional
+
+import simpy
+from openenv_core.env_server import Environment
+
+from ..models import JobT, JSSPAction, JSSPObservation, MachineObservation, ReadyOperationObservation
+
+# Example of JSSP initial jobs
+# Each tuple is a (machine_index, processing_time)
+#
+# FT06: list[JobT] = [
+#     [(2, 1), (0, 3), (1, 6), (3, 7), (5, 3), (4, 6)],
+#     [(1, 8), (2, 5), (4, 10), (5, 10), (0, 10), (3, 4)],
+#     [(2, 5), (3, 4), (5, 8), (0, 9), (1, 1), (4, 7)],
+#     [(1, 5), (0, 5), (2, 5), (3, 3), (4, 8), (5, 9)],
+#     [(2, 9), (1, 3), (4, 5), (5, 4), (0, 3), (3, 1)],
+#     [(1, 3), (3, 3), (5, 9), (0, 10), (4, 4), (2, 1)],
+# ]
+
+PENALTY = 100
+
+
+class JSSPEnvironment(Environment):
+    def __init__(self, jobs: list[JobT]):
+        super().__init__()
+        self.init_jobs = jobs
+        self.reset()
+
+    def reset(self) -> JSSPObservation:
+        """Reset the environment to initial state."""
+        self.episode_id = str(uuid.uuid4())
+        self.step_count = 0
+        self.jobs = deepcopy(self.init_jobs)
+        self.nb_machines = max(max(machine for machine, _ in job) for job in self.jobs) + 1
+
+        # SimPy environment for time tracking
+        self.env = simpy.Environment()
+
+        # Track which operation index each job is currently on
+        self.job_progress = [0] * len(self.jobs)
+
+        # Track machine states
+        self.machine_busy_until: list[Optional[int]] = [None] * self.nb_machines
+        self.machine_current_job: list[Optional[int]] = [None] * self.nb_machines
+
+        # Track completed jobs
+        self.completed_jobs = 0
+
+        return self.state
+
+    def _get_ready_operations(self) -> list[ReadyOperationObservation]:
+        """Get all operations that are ready to be scheduled now."""
+        ready = []
+        for job_id in range(len(self.jobs)):
+            # Skip if job is complete
+            if self.job_progress[job_id] >= len(self.jobs[job_id]):
+                continue
+
+            # Get next operation for this job
+            machine_id, duration = self.jobs[job_id][self.job_progress[job_id]]
+
+            # Check if machine is available
+            busy_until = self.machine_busy_until[machine_id]
+            if busy_until is None or busy_until <= self.env.now:
+                remaining = len(self.jobs[job_id]) - self.job_progress[job_id]
+                ready.append(
+                    ReadyOperationObservation(
+                        job_id=job_id,
+                        machine_id=machine_id,
+                        duration=duration,
+                        remaining_ops=remaining,
+                    )
+                )
+
+        return ready
+
+    def _at_decision_step(self) -> bool:
+        """Check if we're at a decision step (at least one job can be scheduled)."""
+        return len(self._get_ready_operations()) > 0
+
+    def _validate_action(self, action: JSSPAction) -> bool:
+        """Validate that an action is legal."""
+        scheduled_machines = set()
+
+        for job_id in action.job_ids:
+            # Check job ID is valid
+            if job_id < 0 or job_id >= len(self.jobs):
+                return False
+
+            # Check job is not already complete
+            if self.job_progress[job_id] >= len(self.jobs[job_id]):
+                return False
+
+            # Get the machine needed for this job's next operation
+            machine_id, _ = self.jobs[job_id][self.job_progress[job_id]]
+
+            # Check machine is available now
+            busy_until = self.machine_busy_until[machine_id]
+            if busy_until is not None and busy_until > self.env.now:
+                return False
+
+            # Check we're not scheduling two jobs on the same machine
+            if machine_id in scheduled_machines:
+                return False
+
+            scheduled_machines.add(machine_id)
+
+        return True
+
+    def _schedule_jobs(self, job_ids: list[int]):
+        """Schedule the given jobs on their respective machines."""
+        for job_id in job_ids:
+            machine_id, duration = self.jobs[job_id][self.job_progress[job_id]]
+
+            # Update machine state
+            self.machine_busy_until[machine_id] = int(self.env.now) + duration
+            self.machine_current_job[machine_id] = job_id
+
+    def _advance_to_decision_step(self):
+        """Advance simulation time until the next decision step."""
+        while True:
+            # Stop if we're at a decision step
+            if self._at_decision_step():
+                break
+
+            # Stop if all jobs are complete
+            if self.completed_jobs >= len(self.jobs):
+                break
+
+            # Find the next time when a machine becomes free
+            future_times = [t for t in self.machine_busy_until if t is not None and t > self.env.now]
+
+            if not future_times:
+                # No machines will become free, but not all jobs complete
+                # This shouldn't happen in a valid problem
+                break
+
+            next_time = min(future_times)
+
+            # Advance time to when the next machine becomes free
+            self.env.run(until=next_time)
+
+            # Process completed operations and clear machine state
+            for i in range(self.nb_machines):
+                if self.machine_busy_until[i] is not None and self.machine_busy_until[i] <= self.env.now:
+                    # Machine finished processing - advance the job's progress
+                    job_id = self.machine_current_job[i]
+                    if job_id is not None:
+                        self.job_progress[job_id] += 1
+
+                        # Check if job is now complete
+                        if self.job_progress[job_id] >= len(self.jobs[job_id]):
+                            self.completed_jobs += 1
+
+                    # Clear machine state
+                    self.machine_busy_until[i] = None
+                    self.machine_current_job[i] = None
+
+    def step(self, action: JSSPAction) -> JSSPObservation:
+        """Process an action and advance simulation until next decision step.
+
+        Returns observation with reward = -(elapsed time) for valid actions,
+        or reward = -PENALTY for invalid actions (without updating state).
+        """
+        start_time = self.env.now
+
+        # Validate action
+        if not self._validate_action(action):
+            # Invalid action - return current state with penalty
+            obs = self.state
+            obs.reward = -PENALTY
+            return obs
+
+        # Schedule the jobs
+        self._schedule_jobs(action.job_ids)
+
+        # Advance simulation to next decision step
+        self._advance_to_decision_step()
+
+        # Calculate reward as negative time elapsed
+        time_elapsed = self.env.now - start_time
+        reward = -time_elapsed
+
+        # Increment step counter
+        self.step_count = int(self.env.now)
+
+        # Return observation with reward
+        obs = self.state
+        obs.reward = reward
+
+        return obs
+
+    @property
+    def state(self) -> JSSPObservation:
+        """Get the current state of the environment, without the reward."""
+        machines = [
+            MachineObservation(
+                machine_id=i,
+                busy_until=self.machine_busy_until[i],
+                current_job_id=self.machine_current_job[i],
+            )
+            for i in range(self.nb_machines)
+        ]
+
+        ready_ops = self._get_ready_operations()
+
+        return JSSPObservation(
+            done=self.completed_jobs >= len(self.jobs),
+            episode_id=self.episode_id,
+            step_count=self.step_count,
+            machines=machines,
+            ready_operations=ready_ops,
+            completed_jobs=self.completed_jobs,
+            total_jobs=len(self.jobs),
+            reward=0.0,  # Default, overwritten in step()
+        )

src/jssp_openenv/solver.py

@@ -0,0 +1,44 @@
+from .client import JSSPEnvClient
+from .models import ScheduledEvent
+from .policy import JSSPEnvPolicy
+
+
+def solve_jssp(
+    env_client: JSSPEnvClient, policy: JSSPEnvPolicy, max_steps: int, verbose: bool = False
+) -> tuple[int, list[ScheduledEvent]]:
+    """Solve a single JSSP instance using the given policy."""
+    result = env_client.reset()
+    obs = result.observation
+    scheduled_events: list[ScheduledEvent] = []
+
+    while not result.done:
+        if verbose:
+            print(f"Step {obs.step_count}: {obs.ready_operations}")
+        action = policy.act(obs)
+        if verbose:
+            print(f"Action: {action}")
+
+        # Record scheduled events
+        if action.job_ids:
+            for job_id in action.job_ids:
+                operation = next((op for op in obs.ready_operations if op.job_id == job_id), None)
+                assert operation is not None
+                event = ScheduledEvent(
+                    job_id=job_id,
+                    machine_id=operation.machine_id,
+                    start_time=obs.step_count,
+                    end_time=obs.step_count + operation.duration,
+                )
+                scheduled_events.append(event)
+
+        # Execute action
+        result = env_client.step(action)
+        obs = result.observation
+
+        # Safety check to avoid infinite loops
+        if obs.step_count >= max_steps:
+            print(f"\nWARNING: Exceeded max steps ({max_steps}), terminating")
+            break
+
+    # Extract makespan
+    return obs.step_count, scheduled_events

tests/test_models.py

@@ -0,0 +1,17 @@
+import pytest
+
+from jssp_openenv.models import parse_job_ids
+
+
+def test_parse_job_ids():
+    assert parse_job_ids("1,2,3") == [1, 2, 3]
+    assert parse_job_ids("3,2,1") == [3, 2, 1]
+    assert parse_job_ids("") == []
+    assert parse_job_ids(",") == []
+    assert parse_job_ids("0,") == [0]
+
+    with pytest.raises(ValueError):
+        parse_job_ids("1,2,3,a")
+
+    with pytest.raises(ValueError):
+        parse_job_ids("0.1")