# Dealing with Preemption

If your command supports checkpointing, you can try to trigger a checkpoint when your job is about to be preempted.
When using the `ShellCommandProcessor`, be sure to configure your shell to pass the signal to your python process.

```shell
# instead of:
python main.py

# run this:
set -o monitor
python main.py
```

Your process can catch the signal and, for example, trigger saving a checkpoint.

There's an important distinction here between Manifold and AzureML Compute:

- Manifold sends us SIGTERM and then stops the container after a timeout.

- On AzureML Compute, the azure VM announces a preemption, but then Azure may
  reconsider this decision and not preempt the VM after all. If this is too disruptive,
  and you don't have any checkpointing logic anyway, set
  `JOBQ_DISABLE_SCHEDULED_EVENTS=1`.

  The standard behavior is for jobq to send SIGTERM, wait for your command to
  finish, ignore its exit code, and reschedule the task.

If you're *not* using `ShellCommandProcessor`, you likely don't care that much about saving state when preempted.
One thing you can consider is that the `ProcessPool` installs a signal handler that passes the signal to all subprocesses.
Your subprocesses can then install their own signal handlers as they please.


## Preemption walkthrough

If you want your code to be preemptible, you can start with the following sample script which is called dummy_task.py:

```python
import signal
import time


def handle_event(
    event: signal.Signals,
    stack_frame,
):
    """Handle preemption signals."""
    print(f"CHECKPOINTING Received signal: {event}")


def main():
    print("Starting dummy task, will sleep 10 seconds.")
    signal.signal(signal.SIGTERM, handle_event)

    # Simulate a long-running task
    time.sleep(10)
    print("Finished dummy task.")


if __name__ == "__main__":
    main()
```

You can then queue this task:

```python
from ai4s.jobq import JobQ, batch_enqueue
from azure.identity import AzureDefaultCredential

your_queue_name = "demo-preemption"
your_storage_account = "<YOUR-STORAGE-ACCOUNT>"
nb_tasks_to_queue = 40


all_commands_to_run = [
    "set -o monitor && python ./dummy_task.py"
] * nb_tasks_to_queue

# Enqueue the commands to the JobQ queue.
async with JobQ.from_storage_queue(
    your_queue_name, storage_account=your_storage_account, credential=AzureDefaultCredential()
) as queue:
    await batch_enqueue(
        queue,
        all_commands_to_run,
    )
```

You can then start a local ai4s-jobq worker to process the tasks:
Note that this example uses `--num-workers 2` to start two workers on the same node, so two tasks run in parallel.

```shell
ai4s-jobq <YOUR-STORAGE-ACCOUNT>/demo-preemption worker --num-workers 2 --heartbeat --max-consecutive-failures 5 --time-limit 1d
```

To simulate a preemption, press Ctrl-C in the terminal where the worker is running. Look at the output, you can see that the CHECKPOINTING message is printed when the SIGTERM signal is received.

**Best practices**: Your sigterm-handles should be quick, and should not raise exceptions. It should not call sys.exit() or similar. It should ideally just signal to the main thread to checkpoint and exit.

## SIGKILL escalation for unresponsive processes

When a worker sends SIGTERM to a task subprocess (for example, during preemption or
graceful shutdown), the subprocess is expected to clean up and exit. If the process
ignores or mishandles SIGTERM and keeps running, the worker escalates to SIGKILL
after a configurable timeout.

Each subprocess runs in its own process group (`start_new_session=True`), so the
SIGKILL is sent to the entire process group. This ensures that any child processes
spawned by the task are also terminated, preventing orphaned processes from
consuming resources after the worker moves on.

| Environment variable | Default | Description |
|---|---|---|
| `JOBQ_KILL_TIMEOUT` | `600` (10 minutes) | Seconds to wait after SIGTERM before sending SIGKILL to the process group. Set to `0` to disable SIGKILL escalation (not recommended). |

For short-lived test scenarios, you can set a lower value:

```shell
JOBQ_KILL_TIMEOUT=30 ai4s-jobq myaccount/myqueue worker --num-workers 2
```
