From 0ab4c5cd3ac425831d9cbd82094bba78d05e43f1 Mon Sep 17 00:00:00 2001 From: Christian Versloot Date: Tue, 2 May 2023 13:11:05 +0200 Subject: [PATCH] =?UTF-8?q?Generate=20unique=20SageMaker=20training=20job?= =?UTF-8?q?=20name=20based=20on=20pipeline=20and=20ste=E2=80=A6=20(#1505)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Generate unique SageMaker training job name based on pipeline and step name * Client import * Reformatted based on format.sh * Replace random string generation with string_utils.random_str * Revert "Reformatted based on format.sh" This reverts commit 98ca3f1cac3ae94e39537e272d20d7ff190e675a. * Ruff and black changes --------- Co-authored-by: Alex Strick van Linschoten --- .../step_operators/sagemaker_step_operator.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py b/src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py index 02d07af7cc3..3e3355184b2 100644 --- a/src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py +++ b/src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py @@ -17,6 +17,7 @@ import sagemaker +from zenml.client import Client from zenml.config.build_configuration import BuildConfiguration from zenml.enums import StackComponentType from zenml.integrations.aws.flavors.sagemaker_step_operator_flavor import ( @@ -26,6 +27,7 @@ from zenml.logger import get_logger from zenml.stack import Stack, StackValidator from zenml.step_operators import BaseStepOperator +from zenml.utils.string_utils import random_str if TYPE_CHECKING: from zenml.config.base_settings import BaseSettings @@ -179,9 +181,16 @@ def launch( image_name, self.config.role, **estimator_args ) + # SageMaker allows 63 characters at maximum for job name - ZenML uses 60 for safety margin. + step_name = Client().get_run_step(info.step_run_id).name + training_job_name = f"{info.pipeline.name}-{step_name}"[:55] + suffix = random_str(4) + unique_training_job_name = f"{training_job_name}-{suffix}" + # Sagemaker doesn't allow any underscores in job/experiment/trial names - job_name = f"{info.run_name}-{info.pipeline_step_name}" - job_name = job_name.replace("_", "-") + sanitized_training_job_name = unique_training_job_name.replace( + "_", "-" + ) # Construct training input object, if necessary inputs = None @@ -201,12 +210,12 @@ def launch( if settings.experiment_name: experiment_config = { "ExperimentName": settings.experiment_name, - "TrialName": job_name, + "TrialName": sanitized_training_job_name, } estimator.fit( wait=True, inputs=inputs, experiment_config=experiment_config, - job_name=job_name, + job_name=sanitized_training_job_name, )