[New Resource]: AWS SageMaker MLFlow Tracking Server

[philipgebus]

Description

AWS SageMaker has recently released experiments tracking via a MLFlow Tracking Server.

Requested Resource(s) and/or Data Source(s)

aws_sagemaker_mlflow_tracking_server

Potential Terraform Configuration

resource "aws_sagemaker_mlflow_tracking_server" "this" {
  tracking_server_name = "mlflow-dev"
  artifact_store_uri = "s3://bucket/prefix/"
  
  tracking_server_size = "Small|Medium|Large"
  mlflow_version = "2.13.2"
  
  role_arn = "arn:aws:iam::000000000000:role/tracking-server-role"
  
  automatic_model_registration = true
  weekly_maintenance_window_start = "TUE:03:30"

  tags = {
    sagemaker:user-profile-arn           = "arn:aws:sagemaker:eu-central-1:0000000000:user-profile/d-***/johndoe",
    sagemaker:domain-arn     = aws_sagemaker_domain.mlflow.arn
  }
}

References

AWS blog post:
https://aws.amazon.com/blogs/aws/manage-ml-and-generative-ai-experiments-using-amazon-sagemaker-with-mlflow

Code reference:
https://pkg.go.dev/github.com/aws/aws-sdk-go-v2/service/sagemaker#Client.CreateMlflowTrackingServer
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_mlflow_tracking_server.html

Would you like to implement a fix?

No

[github-actions]

Community Note

Voting for Prioritization

• Please vote on this issue by adding a 👍 reaction to the original post to help the community and maintainers prioritize this request.
• Please see our prioritization guide for information on how we prioritize.
• Please do not leave "+1" or other comments that do not add relevant new information or questions, they generate extra noise for issue followers and do not help prioritize the request.

Volunteering to Work on This Issue

• If you are interested in working on this issue, please leave a comment.
• If this would be your first contribution, please review the contribution guide.

[jmeisele]

How do you think this would this work in a centralized model registry located in another AWS Account? Would the MLFlow server just be located in the same AWS Account as the model registry, then the MLFlow tracking server be referenced in each related sagemaker domain(s) in dev, qa, prod?

[philipgebus]

How do you think this would this work in a centralized model registry located in another AWS Account? Would the MLFlow server just be located in the same AWS Account as the model registry, then the MLFlow tracking server be referenced in each related sagemaker domain(s) in dev, qa, prod?

In my preferred solution the MLFlow tracking server operates within a centralized model registry AWS account and is accessed as mentioned by SageMaker domains located in various AWS accounts. However, it seems that the SageMaker MLFlow tracking server does not (yet?) support IAM resource policies, which would facilitate convenient cross-account access. Hence, one would either need to utilize sts:AssumeRole (SageMaker job role assumes a mlflow role in the central model registry account before logging training runs) or come up with a custom synchronization approach in order to work in a cross-account setup.

[DrFaust92]

Thinking of working on this, talking to provider team on how to tackle it as new resources require go sdk v2 and its not yet supported with sagemaker in the provider yet AFAIK and need to understand how much work this is 😸

[jorge-trivilin]

This would be good. Set up a new tracking server for every new project, experiment or branch setup. ;)

[jmeisele]

Good Morning ☕ @deepakbshetty and @DrFaust92 have either of you had a chance to take a look at this new resource for sagemaker? Looks like the go sdk v2 came out with mlflow functionality as part of release 2024-06-18.

I found the:

• Start
• Stop
• Update

Go commands if that is any help 🫶

[DrFaust92]

jmeisele Started work on this, still missing a few things but a I have a working resource now

[github-actions]

Warning

This issue has been closed, meaning that any additional comments are hard for our team to see. Please assume that the maintainers will not see them.

Ongoing conversations amongst community members are welcome, however, the issue will be locked after 30 days. Moving conversations to another venue, such as the AWS Provider forum, is recommended. If you have additional concerns, please open a new issue, referencing this one where needed.

[github-actions]

This functionality has been released in v5.73.0 of the Terraform AWS Provider. Please see the Terraform documentation on provider versioning or reach out if you need any assistance upgrading.

For further feature requests or bug reports with this functionality, please create a new GitHub issue following the template. Thank you!