Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Generate unique SageMaker training job name based on pipeline and ste… #1505

Merged
merged 7 commits into from
May 2, 2023
Merged

Generate unique SageMaker training job name based on pipeline and ste… #1505

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
3d33587
Generate unique SageMaker training job name based on pipeline and ste…
christianversloot Apr 24, 2023
a182855
Client import
christianversloot Apr 24, 2023
98ca3f1
Reformatted based on format.sh
christianversloot Apr 24, 2023
4e1fe24
Replace random string generation with string_utils.random_str
christianversloot Apr 25, 2023
be06686
Merge branch 'develop' into bugfix/sagemaker-training-job-name
strickvl Apr 28, 2023
af51e83
Revert "Reformatted based on format.sh"
christianversloot May 2, 2023
5daed8b
Ruff and black changes
christianversloot May 2, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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,
)