Platform For AI:Submit a training job

Step 1: Prepare your training script

Write your training script locally. PAI provisions the cloud environment and runs it. The script must follow these conventions to integrate with the PAI runtime.

Environment variables available in the training container

PAI injects the following environment variables into every training container. Reference them in your training script or startup command.

Load hyperparameters

When you set hyperparameters on the Estimator, PAI writes the values to {PAI_CONFIG_DIR}/hyperparameters.json (default path: /ml/input/config/hyperparameters.json).

Read hyperparameters from this file:

import json
import os

def load_hyperparameters():
    hps_path = os.path.join(os.environ["PAI_CONFIG_DIR"], "hyperparameters.json")
    with open(hps_path) as f:
        return json.load(f)

For example, hyperparameters={"batch_size": 32, "learning_rate": 0.01} produces:

{
  "batch_size": "32",
  "learning-rate": "0.01"
}

Alternatively, use PAI_USER_ARGS to forward all hyperparameters to your script as command-line arguments:

est = Estimator(
    command="python train.py $PAI_USER_ARGS",
    hyperparameters={
        "epochs": 10,
        "batch-size": 32,
        "learning-rate": 0.001,
    },
    # ...
)

Parse them in the training script using argparse.

For hyperparameters={"epochs": 10, "batch-size": 32, "train.learning_rate": 0.001}, PAI generates:

PAI_HPS_EPOCHS=10
PAI_HPS_BATCH_SIZE=32
PAI_HPS_TRAIN_LEARNING_RATE=0.001

Load input data

Pass input data via the inputs parameter of Estimator.fit(). Each entry is a key-value pair where the key is the channel name and the value is an OSS or File Storage NAS (NAS) path:

est.fit(
    inputs={
        "train": "oss://<YourOssBucket>/train/data/train.csv",
        "test": "oss://<YourOssBucket>/test/data/",
    }
)

PAI mounts the data at /ml/input/data/{channel_name}/. Read the mount path using PAI_INPUT_{CHANNEL_NAME}:

train_data = os.environ["PAI_INPUT_TRAIN"]   # /ml/input/data/train/
test_data  = os.environ["PAI_INPUT_TEST"]    # /ml/input/data/test/

If the OSS path ends with /, the environment variable points to a directory. If it ends with a filename, it points to that file.

Save the output model

Save the trained model to the path specified by PAI_OUTPUT_MODEL (default: /ml/output/model/). PAI automatically uploads anything written there to OSS.

import os

def save_model(model, model_dir=None):
    if model_dir is None:
        model_dir = os.environ["PAI_OUTPUT_MODEL"]
    # Write your model files to model_dir.
    # Example: torch.save(model.state_dict(), os.path.join(model_dir, "model.pth"))

Full script example

import json
import os


def load_hyperparameters():
    hps_path = os.path.join(os.environ["PAI_CONFIG_DIR"], "hyperparameters.json")
    with open(hps_path) as f:
        return json.load(f)


def train(hps, train_data, test_data):
    """Add your model training code here."""
    pass


def save_model(model):
    output_path = os.environ["PAI_OUTPUT_MODEL"]
    # Write the model to output_path.
    pass


def run():
    # 1. Load hyperparameters.
    hps = load_hyperparameters()

    # 2. Load input data from mounted paths.
    train_data = os.environ["PAI_INPUT_TRAIN"]
    test_data  = os.environ["PAI_INPUT_TEST"]

    model = train(hps, train_data, test_data)

    # 3. Save the output model.
    save_model(model)


if __name__ == "__main__":
    run()

Add third-party dependencies

If your script requires Python packages not included in the training image, add a requirements.txt file in the same directory as your script. PAI installs the listed packages before running the script.

Organize your code directory as follows:

train_src/
├── requirements.txt   # Third-party dependencies
├── train.py           # Entry point; run with: python train.py
└── utils.py

Set source_dir="train_src" on the Estimator to package and upload this directory to PAI.

Step 2: Choose a training image

Specify a container image that includes your framework and library dependencies. Use a PAI prebuilt image or a custom image from Container Registry (ACR).

Use a PAI prebuilt image

from pai.image import retrieve, list_images, ImageScope

# List all PAI images for PyTorch.
for image_info in list_images(framework_name="PyTorch"):
    print(image_info)

# Get the TensorFlow 2.3 image for CPU training.
tf_image = retrieve(framework_name="TensorFlow", framework_version="2.3")

# Get the latest TensorFlow image for GPU training.
tf_gpu_image = retrieve(
    framework_name="TensorFlow",
    framework_version="latest",
    accelerator_type="GPU",
)

# Get the PyTorch 1.12 image for GPU training.
torch_gpu_image = retrieve(framework_name="PyTorch", framework_version="1.12",
                           accelerator_type="GPU")

For a list of third-party libraries preinstalled in PAI images, see Public images.

Step 3: Submit the job

Create an Estimator instance and call fit() to submit the job. After submission, PAI prints the job details page URL and streams logs until the job succeeds, fails, or stops.

from pai.estimator import Estimator
from pai.image import retrieve

torch_image_uri = retrieve("PyTorch", framework_version="1.12").image_uri

est = Estimator(
    command="python train.py",
    source_dir="./train_src/",
    image_uri=torch_image_uri,
    instance_type="ecs.c6.large",
    hyperparameters={
        "n_estimators": 500,
        "objective": "reg:squarederror",
        "max_depth": 5,
    },
    base_job_name="example_train_job",
)

est.fit()

# Get the OSS path of the output model.
print(est.model_data())

Key parameters:

• command — The startup command that PAI runs inside the container (for example, python train.py).

• source_dir — Local directory containing your training script and dependencies. PAI packages and uploads its contents. Also accepts an OSS path to a TAR archive (for example, oss://<YourOssBucket>/source.tar.gz). If requirements.txt exists in this directory, dependencies are installed automatically before the job runs.

• image_uri — Container image URI. Use retrieve() to get a PAI prebuilt image URI, or provide a custom ACR image URI.

• instance_type — ECS instance type for training (for example, ecs.c6.large).

• hyperparameters — Key-value pairs passed to the training script. Values are written to hyperparameters.json and also exposed as environment variables.

• base_job_name — Prefix for the job name. The full name follows the format {base_job_name}_{submitted-datetime}.

After fit() returns, call est.model_data() to get the OSS path of the saved model.

Debug locally

Debugging in the cloud can be slow. Run and debug a training job locally by setting instance_type="local". The job runs in a local Docker container.

est = Estimator(
    command="python train.py",
    source_dir="./train_src/",
    image_uri=torch_image_uri,
    instance_type="local",
)

est.fit(
    inputs={
        # OSS data is downloaded and mounted into the container.
        "train": "oss://<BucketName>/path-to-data/",
        # Local paths are mounted directly.
        "test": "/data/to/test/data",
    }
)

print(est.model_data())