Osservabilità dei modelli per i lavori di formazione su SageMaker HyperPod cluster orchestrati da Amazon EKS

Crea un server di MLflow tracciamento SageMaker AI, seguendo le istruzioni in Creare un server di tracciamento utilizzando la AWS CLI.

Assicurati che l'eks-auth:AssumeRoleForPodIdentityautorizzazione esista nel ruolo di esecuzione IAM per SageMaker HyperPod.

Se il eks-pod-identity-agent componente aggiuntivo non è già installato sul cluster EKS, installalo sul cluster EKS.

aws eks create-addon \
    --cluster-name <eks_cluster_name> \
    --addon-name eks-pod-identity-agent \
    --addon-version vx.y.z-eksbuild.1

Crea un trust-relationship.json file per un nuovo ruolo da chiamare per Pod. MLflow APIs

cat >trust-relationship.json <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "AllowEksAuthToAssumeRoleForPodIdentity",
            "Effect": "Allow",
            "Principal": {
                "Service": "pods.eks.amazonaws.com"

            },
            "Action": [
                "sts:AssumeRole",
                "sts:TagSession"
            ]
        }
    ]
}
EOF

Esegui il codice seguente per creare il ruolo e collegare la relazione di fiducia.

aws iam create-role --role-name hyperpod-mlflow-role \
    --assume-role-policy-document file://trust-relationship.json \
    --description "allow pods to emit mlflow metrics and put data in s3"

Crea la seguente politica che concede a Pod l'accesso per chiamare tutte le sagemaker-mlflow operazioni e inserire gli artefatti del modello in S3. L'autorizzazione S3 esiste già all'interno del server di tracciamento, ma se gli artefatti del modello sono troppo grandi, viene effettuata una chiamata diretta a s3 dal codice per caricare gli artefatti. MLflow

cat >hyperpod-mlflow-policy.json <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "sagemaker-mlflow:AccessUI",
                "sagemaker-mlflow:CreateExperiment",
                "sagemaker-mlflow:SearchExperiments",
                "sagemaker-mlflow:GetExperiment",
                "sagemaker-mlflow:GetExperimentByName",
                "sagemaker-mlflow:DeleteExperiment",
                "sagemaker-mlflow:RestoreExperiment",
                "sagemaker-mlflow:UpdateExperiment",
                "sagemaker-mlflow:CreateRun",
                "sagemaker-mlflow:DeleteRun",
                "sagemaker-mlflow:RestoreRun",
                "sagemaker-mlflow:GetRun",
                "sagemaker-mlflow:LogMetric",
                "sagemaker-mlflow:LogBatch",
                "sagemaker-mlflow:LogModel",
                "sagemaker-mlflow:LogInputs",
                "sagemaker-mlflow:SetExperimentTag",
                "sagemaker-mlflow:SetTag",
                "sagemaker-mlflow:DeleteTag",
                "sagemaker-mlflow:LogParam",
                "sagemaker-mlflow:GetMetricHistory",
                "sagemaker-mlflow:SearchRuns",
                "sagemaker-mlflow:ListArtifacts",
                "sagemaker-mlflow:UpdateRun",
                "sagemaker-mlflow:CreateRegisteredModel",
                "sagemaker-mlflow:GetRegisteredModel",
                "sagemaker-mlflow:RenameRegisteredModel",
                "sagemaker-mlflow:UpdateRegisteredModel",
                "sagemaker-mlflow:DeleteRegisteredModel",
                "sagemaker-mlflow:GetLatestModelVersions",
                "sagemaker-mlflow:CreateModelVersion",
                "sagemaker-mlflow:GetModelVersion",
                "sagemaker-mlflow:UpdateModelVersion",
                "sagemaker-mlflow:DeleteModelVersion",
                "sagemaker-mlflow:SearchModelVersions",
                "sagemaker-mlflow:GetDownloadURIForModelVersionArtifacts",
                "sagemaker-mlflow:TransitionModelVersionStage",
                "sagemaker-mlflow:SearchRegisteredModels",
                "sagemaker-mlflow:SetRegisteredModelTag",
                "sagemaker-mlflow:DeleteRegisteredModelTag",
                "sagemaker-mlflow:DeleteModelVersionTag",
                "sagemaker-mlflow:DeleteRegisteredModelAlias",
                "sagemaker-mlflow:SetRegisteredModelAlias",
                "sagemaker-mlflow:GetModelVersionByAlias"
            ],
            "Resource": "arn:aws:sagemaker:us-west-2:111122223333:mlflow-tracking-server/<ml tracking server name>"
        },
        {
            "Effect": "Allow",
            "Action": [
                "s3:PutObject"
            ],
            "Resource": "arn:aws:s3:::<mlflow-s3-bucket_name>"
        }
    ]
}
EOF

Allega la mlflow-metrics-emit-policy politica al documento relativo all'hyperpod-mlflow-roleutilizzo della policy salvato nel passaggio precedente.

aws iam put-role-policy \
  --role-name hyperpod-mlflow-role \
  --policy-name mlflow-metrics-emit-policy \
  --policy-document file://hyperpod-mlflow-policy.json

Crea un account di servizio Kubernetes per consentire a Pod di accedere al server. MLflow

cat >mlflow-service-account.yaml <<EOF
apiVersion: v1
kind: ServiceAccount
metadata:
  name: mlflow-service-account
  namespace: kubeflow
EOF

Esegui il comando seguente da applicare al cluster EKS.

kubectl apply -f mlflow-service-account.yaml

Crea un'associazione di identità Pod.

aws eks create-pod-identity-association \
    --cluster-name EKS_CLUSTER_NAME \
    --role-arn arn:aws:iam::111122223333:role/hyperpod-mlflow-role \
    --namespace kubeflow \
    --service-account mlflow-service-account

Aggiungi le seguenti righe all'inizio dello script di formazione.

import mlflow

# Set the Tracking Server URI using the ARN of the Tracking Server you created
mlflow.set_tracking_uri(os.environ['MLFLOW_TRACKING_ARN'])
# Enable autologging in MLflow
mlflow.autolog()

Crea un'immagine Docker con lo script di formazione e inviala ad Amazon ECR. Ottieni l'ARN del contenitore ECR. Per ulteriori informazioni sulla creazione e il push di un'immagine Docker, consulta Pushing a Docker image nella ECR User Guide.

Aggiungi un account di servizio nei Job Pods di formazione a cui consentire loro di accedere. hyperpod-mlflow-role Ciò consente ai Pods di chiamare. MLflow APIs Esegui il seguente modello di invio di lavori SageMaker HyperPod CLI. Createlo con il nome del file. mlflow-test.yaml

defaults:
 - override hydra/job_logging: stdout

hydra:
 run:
  dir: .
 output_subdir: null

training_cfg:
 entry_script: ./train.py
 script_args: []
 run:
  name: test-job-with-mlflow # Current run name
  nodes: 2 # Number of nodes to use for current training
  # ntasks_per_node: 1 # Number of devices to use per node
cluster:
 cluster_type: k8s # currently k8s only
 instance_type: ml.c5.2xlarge
 cluster_config:
  # name of service account associated with the namespace
  service_account_name: mlflow-service-account
  # persistent volume, usually used to mount FSx
  persistent_volume_claims: null
  namespace: kubeflow
  # required node affinity to select nodes with SageMaker HyperPod
  # labels and passed health check if burn-in enabled
  label_selector:
      required:
          sagemaker.amazonaws.com/node-health-status:
              - Schedulable
      preferred:
          sagemaker.amazonaws.com/deep-health-check-status:
              - Passed
      weights:
          - 100
  pullPolicy: IfNotPresent # policy to pull container, can be Always, IfNotPresent and Never
  restartPolicy: OnFailure # restart policy

base_results_dir: ./result # Location to store the results, checkpoints and logs.
container: 111122223333.dkr.ecr.us-west-2.amazonaws.com/tag # container to use

env_vars:
 NCCL_DEBUG: INFO # Logging level for NCCL. Set to "INFO" for debug information
 MLFLOW_TRACKING_ARN: arn:aws:sagemaker:us-west-2:11112223333:mlflow-tracking-server/tracking-server-name

Avviate il lavoro utilizzando il file YAML come segue.

hyperpod start-job --config-file /path/to/mlflow-test.yaml

Genera un URL prefirmato per il MLflow server di tracciamento. Puoi aprire il link sul tuo browser e iniziare a monitorare il tuo lavoro di formazione.

aws sagemaker create-presigned-mlflow-tracking-server-url \                          
    --tracking-server-name "tracking-server-name" \
    --session-expiration-duration-in-seconds 1800 \
    --expires-in-seconds 300 \
    --region region