[Infrastructure] AWS & SageMaker Enhancements (Spot Instances, Cost Safeguards & Security)

[igor-soldev]

Hello! 👋

I recently ran the AWS/SageMaker templates through InfraScan — an Open Source tool my team is developing to audit infrastructure-as-code against cloud best practices.

The codebase scored a solid "B" (91.3%), but the scan highlighted a few specific areas where we could improve cost efficiency (crucial for MLOps) and security posture:

1. SageMaker Spot Instances (Cost Optimization)
The aws_sagemaker_notebook_instance configuration (aws/terraform/sagemaker-notebook/sagemaker.tf) currently defaults to On-Demand instances.

• Impact: ML compute is expensive.
• Suggestion: Since notebook instances are often used for experimentation, adding a variable/toggle to enable specific Lifecycle Configurations that leverage Spot Instances could save users ~70% on compute costs.

2. Missing AWS Budget (Cost Safety)
The AWS templates don't appear to provision an aws_budgets_budget by default.

• Impact: MLOps workloads can spiral quickly if a large job is left running. Adding a default budget alert is a great safety net for new users trying out the stack.

3. SageMaker IMDSv2 (Security)
The notebook instances currently allow IMDSv1 (Instance Metadata Service Version 1).

• Impact: In shared ML environments where users execute arbitrary code (Jupyter), IMDSv1 leaves the instance credential endpoint vulnerable to SSRF attacks.
• Suggestion: Enforcing IMDSv2 (http_tokens = "required") is a standard hardening step for SageMaker.

If you'd like to see the specific lines of code, the full interactive report is available here:

👉 View Full InfraScan Report for metaflow-tools

(Full disclosure: the link above is generated by our tool, but I manually reviewed the findings to make sure they are relevant to the Metaflow architecture).

If you're open to it, I'd be happy to submit a PR to add the AWS Budget resource and harden the SageMaker configuration. Let me know what you think!