Metadata-Version: 2.5
Name: toolkit-for-ray-on-sagemaker-ai
Version: 1.0.4
Summary: Amazon SageMaker Ray library
License: Apache 2.0
License-File: LICENSE.txt
Requires-Python: >=3.11
Requires-Dist: aiohttp
Requires-Dist: boto3
Requires-Dist: pyyaml>=6.0
Requires-Dist: ray[default]>=2.0.0
Requires-Dist: requests>=2.20.0
Description-Content-Type: text/markdown

# Toolkit for Ray on Amazon SageMaker AI

Amazon SageMaker Ray library — utilities that make running Ray workloads on
Amazon SageMaker HyperPod simpler and more reliable.

## Overview

The `toolkit-for-ray-on-sagemaker-ai` library adds SageMaker HyperPod
capabilities to your Ray workflows using standard Ray APIs:

- **Authenticated job submission** — submit jobs to secured HyperPod Ray
  clusters through Ray's native `JobSubmissionClient` and `ray job` CLI, with
  authentication handled transparently.
- **Hung job detection** — define log-based rules that detect when a Ray Train
  job stops making progress, and optionally cancel it automatically.
- **JumpStart model loading** — download Amazon SageMaker JumpStart model
  artifacts onto Ray worker nodes for Ray Serve LLM deployments.

## Installation

```bash
pip install toolkit-for-ray-on-sagemaker-ai
```

Requires Python 3.11+ and Ray 2.0.0+ (`ray[default]`).

## Job Submission

The library registers a `sagemaker_ray://` address scheme for Ray's job
submission APIs. When you use this scheme, the library acquires credentials and
establishes an authenticated session for you — no custom client classes and no
explicit imports are required.

### Prerequisites

- A SageMaker HyperPod cluster (orchestrated by Amazon EKS) running a Ray
  cluster with the public Ray dashboard endpoint enabled.
- `kubectl` configured for the cluster:
  `aws eks update-kubeconfig --name <cluster>`. The library reads your
  kubeconfig (`~/.kube/config` or `$KUBECONFIG`).
- The AWS CLI installed and AWS credentials configured — the EKS kubeconfig
  uses `aws eks get-token` to obtain a token.
- Job submission is restricted to the cluster's owner.

### Python SDK

```python
from ray.job_submission import JobSubmissionClient

# Address form: sagemaker_ray://<cluster-name>/<namespace>
client = JobSubmissionClient("sagemaker_ray://my-cluster/team-a")

job_id = client.submit_job(entrypoint="python train.py")
print(client.get_job_status(job_id))
print(client.get_job_logs(job_id))
```

### CLI

```bash
ray job submit --address "sagemaker_ray://my-cluster/team-a" -- python train.py
ray job status --address "sagemaker_ray://my-cluster/team-a" <job_id>
ray job logs   --address "sagemaker_ray://my-cluster/team-a" <job_id>
```

### Address formats

- **Cluster name**: `sagemaker_ray://<cluster-name>/<namespace>`. The namespace
  defaults to `default` if omitted: `sagemaker_ray://my-cluster`.
- **Dashboard URL**: `sagemaker_ray://<dashboard-host>` when you already have
  the fully qualified dashboard hostname.

## Hung Job Detection

`SageMakerLogMonitoring` watches your training job's stdout for expected log
patterns. If an expected pattern stops appearing within a configured window (or
an error pattern appears), the job is flagged as hung and can be cancelled
automatically. A rule can also define a `stop_pattern` that deactivates it once
the job reaches a known terminal state (e.g. training completes), so normal
shutdown isn't mistaken for a hang.

```python
from toolkit_for_ray_on_sagemaker_ai.log_monitoring import SageMakerLogMonitoring, LogMonitorConfig

def train_func():
    SageMakerLogMonitoring(config=LogMonitorConfig(
        enabled=True,
        rules=[
            {
                "name": "training_progress",
                "type": "log_pattern",
                "enabled": True,
                "log_pattern": r"(Epoch|Step|Iteration) \d+",
                "timeout_minutes": 10,
                "start_timeout_minutes": 30,
                "stop_pattern": "Training complete",
                "metric_evaluation_data_points": 3,
                "fault_on_match": False,
            },
            {
                "name": "oom_detection",
                "type": "log_pattern",
                "enabled": True,
                "log_pattern": "CUDA out of memory|OutOfMemoryError|OOM",
                "fault_on_match": True,
            },
        ],
        action="cancel",
    )).start()

    # ... your training loop ...

# To opt out of hang detection entirely:
SageMakerLogMonitoring(config=LogMonitorConfig(enabled=False)).start()
```

### Rule fields

| Field | Type | Description |
|---|---|---|
| `name` | str | Human-readable identifier for the rule. |
| `type` | str | Signal type. Use `"log_pattern"` for log-based detection. |
| `enabled` | bool | Whether this individual rule is active. |
| `log_pattern` | str | Regex to match in stdout (RE2 syntax, max 256 chars). |
| `timeout_minutes` | int | How long the pattern may be absent before a hang is declared. |
| `start_timeout_minutes` | int | Max time from job start to first match (allows for startup/model loading). |
| `stop_pattern` | str | Optional regex that deactivates the rule when matched (e.g. `"Training complete"`). |
| `metric_evaluation_data_points` | int | Number of consecutive evaluations a rule must fail before a hang is declared. Default `1`; use a higher value to require sustained violations and reduce false positives. |
| `fault_on_match` | bool | If `True`, declares a hang immediately when the pattern matches (use for error patterns like OOM). |

### `LogMonitorConfig`

| Field | Type | Default | Description |
|---|---|---|---|
| `enabled` | bool | `True` | Enable/disable hang detection for the job. |
| `rules` | list[dict] | `[]` | Detection rule definitions (see above). |
| `action` | str | `"notify"` | `"notify"` emits a detection event only; `"cancel"` terminates the hung training process. |

> Hung job detection relies on the monitoring capability provided by
> SageMaker HyperPod. On clusters where it is not available, calls are a
> no-op and your training job runs normally.

## JumpStart Model Loading

`JumpStartModelLoaderCallback` downloads SageMaker JumpStart model artifacts to
each Ray worker node before the serving engine initializes, for use with Ray
Serve LLM deployments.

```python
from toolkit_for_ray_on_sagemaker_ai.jumpstart import JumpStartModelLoaderCallback

callback_config = CallbackConfig(
    callback_class=JumpStartModelLoaderCallback,
    callback_kwargs={
        "jumpstart_model_id": "meta-textgeneration-llama-3-1-8b-instruct",
        "region": "us-east-1",
        "accept_eula": True,
    },
)
```

## Troubleshooting

- **"Access denied … Only the cluster creator can submit jobs."** — Job
  submission is restricted to the cluster owner. Use credentials for the
  identity that created the cluster.
- **"Kubeconfig not found" / authentication failed** — Run
  `aws eks update-kubeconfig --name <cluster>` and confirm your AWS credentials
  are valid.
- **"The public endpoint may not be enabled for this cluster."** — Ensure the
  cluster's public Ray dashboard endpoint is enabled.
- **Resource not found** — Verify the cluster name and namespace in the address.

## Requirements

- Python >= 3.11
- Ray >= 2.0.0 (`ray[default]`)
- AWS CLI installed and credentials configured

## Dependencies

`requests`, `pyyaml`, `ray[default]`, `aiohttp`, `boto3`

## License

Apache License 2.0. See `LICENSE.txt` in the package for details.
