Submit Slurm Training Jobs Quickly

Step 1: Prepare the training script and dataset

Prepare the following training dataset and script, and then upload them to your OSS bucket. For more information, see Console Quick Start.

• Training dataset: mnist.npz.

• Training script: mnist_train.py. The following code provides an example:

  import numpy as np
  import tensorflow as tf
  from tensorflow.keras import layers, models
  
  # Load the MNIST dataset from a local file.
  with np.load('/mnt/test/mnist.npz', allow_pickle=True) as data:
      train_images = data['x_train']
      train_labels = data['y_train']
      test_images = data['x_test']
      test_labels = data['y_test']
  
  # Normalize the dataset.
  train_images, test_images = train_images / 255.0, test_images / 255.0
  
  # Add a channel dimension.
  train_images = train_images[..., tf.newaxis]
  test_images = test_images[..., tf.newaxis]
  
  # Build the model.
  model = models.Sequential([
      layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
      layers.MaxPooling2D((2, 2)),
      layers.Conv2D(64, (3, 3), activation='relu'),
      layers.MaxPooling2D((2, 2)),
      layers.Conv2D(64, (3, 3), activation='relu'),
      layers.Flatten(),
      layers.Dense(64, activation='relu'),
      layers.Dense(10)
  ])
  
  # Compile the model.
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['accuracy'])
  
  # Create a TensorBoard callback.
  tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir="/mnt/test/mnist_logs", histogram_freq=1)
  
  # Train the model.
  model.fit(train_images, train_labels, epochs=5,
            validation_data=(test_images, test_labels),
            callbacks=[tensorboard_callback])

Step 2: Submit the training job

1. Go to the Create Job page.

   1. Log on to the PAI console. select the target region and workspace at the top of the page, and then click Deep Learning Containers (DLC).

   2. On the Deep Learning Containers (DLC) page, click Create Job.

2. On the Create Job page, configure the following key parameters. For more information about other parameters, see Create a training job.

   Parameter		Description
   Basic Information	Job Name	Enter a custom job name, such as mnist_cnn_tb.
   Environment Information	Image Configuration	Use a Linux-based runtime image. No special requirements apply. This solution uses the Alibaba Cloud Image > tensorflow-training:2.11-gpu-py39-cu112-ubuntu20.04.
   	Mount storage	If your script file is large, attach it to a specified directory in the container. This job uses OSS as an example to show you how to configure storage mounts: Click OSS, and configure the following parameters: URI: Enter the OSS path where the mnist_train.py script is stored. Mount Path: Enter the path in the DLC container where you want to attach the script file. For example, /mnt/test/.
   	Startup Command	The command that the job runs: Start the command with #!/bin/bash. You can customize SBATCH parameters. However, avoid manually setting resource allocation parameters. The platform automatically allocates resources based on your needs. # Supported. #SBATCH --job-name=mnist_cnn_tb # Job name. #SBATCH --output=/mnt/test/mnist_train_%j.out # Output and error log file name. # Not recommended. Correctness is not guaranteed. #SBATCH --cpus-per-task=<N> # Number of CPU cores per task. #SBATCH --mem=<M> # Memory required, in MB. #SBATCH --mem-per-cpu=<M> # Memory per CPU core. #SBATCH --gres=<resource>:<N> # Specific resources requested, such as GPU (gpu:1). #SBATCH --nodes=<N> # Number of nodes requested. #SBATCH --ntasks=<N> # Total number of tasks. #SBATCH --ntasks-per-node=<N> # Number of tasks per node. #SBATCH --partition=<partition> # Partition to submit the job to. #SBATCH --time=<time> # Maximum run time for the job. #SBATCH --account=<account> # Account to use. This job uses the following start command: #!/bin/bash #SBATCH --job-name=mnist_cnn_tb # Job name. #SBATCH --output=/mnt/test/mnist_train_%j.out # Output and error log file name. # Run the training script. srun /usr/bin/python3 /mnt/test/mnist_train.py
   Resource Information	Source	Select Public Resources. You can also choose other resource types as needed.
   	Framework	Select Slurm.
   	Job Resource	Configure Master and Worker resources as follows: Number of Nodes: Set to 1. Resource Type: Set to ecs.g6.xlarge.
   	Retention Period	Set the retention period for instances after job completion to 1 hour. This lets you enter the container to troubleshoot issues if the job fails. Note Enabling job retention consumes system resources continuously.

3. Click OK.

   The Task Precheck dialog box appears.

   

   The Slurm Sbatch Command is automatically generated by the system to submit the job. After submission, the system attaches the script to /usr/config/user_script and runs it using the selected resource parameters.

4. Click Submit Anyway.

   You are redirected to the Deep Learning Containers (DLC) page, where you can view the submitted training job.