Repository navigation
feat(mlops): Add AthenaStep pipeline step - #6508
Conversation
Add AthenaStep, which takes the Athena StartQueryExecution request as a dict and emits it unchanged as the step's Arguments. The service validates the request at CreatePipeline, so the SDK adds no field vocabulary of its own. Properties follow the GetQueryExecutionOutput shape, so downstream references use the QueryExecution prefix. Integ tests passed against a preprod stage. --- X-AI-Prompt: Port the AthenaStep to the latest master of the PySDK and commit it on feat/athena-step X-AI-Tool: kiro-cli
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
| def __init__( | ||
| self, | ||
| name: str, | ||
| arguments: Dict[str, Any], |
There was a problem hiding this comment.
This requires customers to hand-build the full StartQueryExecution dict with every key and the nesting (ResultConfiguration.OutputLocation, ResultConfiguration.EncryptionConfiguration.EncryptionOption). Check how EMRStep/EMRServerlessStep handle this. They take named params (EMRStepConfig/EMRServerlessJobConfig) and assemble the dict internally and neither validates against the service shape, so named params and server-side validation are independent. Can we match that here: named params (or a small AthenaQueryConfig) build the dict, service still validates, and keep a raw arguments escape hatch for forward-compat?
Customers should not hand-build the nested StartQueryExecution dict. AthenaQueryConfig takes named params and assembles the request, like EMRServerlessJobConfig; the service still validates it at CreatePipeline. A raw arguments dict remains as an escape hatch for fields the config does not cover. --- X-AI-Prompt: Address review asking for named params instead of a raw dict on AthenaStep X-AI-Tool: kiro-cli
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
| query_config: Optional[AthenaQueryConfig] = None, | ||
| arguments: Optional[Dict[str, Any]] = None, |
There was a problem hiding this comment.
We can collapse query_config + arguments into a single union-typed param to make this simpler.
query_config: Union[AthenaQueryConfig, Dict[str, Any]]
That customer do not have to worry about two separate fields and also simplifies our code. User can either pass AthenaQueryConfig or build request on their own if they need any additional properties.
In the doc string, you can callout that AthenaQueryConfig should be the preference, unless it is missing some field.
Take a single query_config that is either an AthenaQueryConfig or a raw StartQueryExecution dict, so customers have one field to set. The docstring prefers AthenaQueryConfig and keeps the dict for fields it does not cover. --- X-AI-Prompt: Collapse query_config and arguments into one union param X-AI-Tool: kiro-cli
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
--- X-AI-Prompt: Add an integ test for the raw dict path X-AI-Tool: kiro-cli
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 AI code review
This PR adds an AthenaStep pipeline step plus an AthenaQueryConfig helper to sagemaker-mlops, loads the botocore athena service model into PropertiesMeta so step.properties is shaped like GetQueryExecutionOutput, adds the ATHENA step type enum, and exports the new classes. The implementation closely follows the existing EMRStep raw-arguments pattern and is well tested (unit + integ). The code is correct and consistent with project conventions; the only issue I found is an incorrect import in the module docstring example.
Reviewed commit 0efd6c2. Automated review; verify before acting.
|
|
||
| from sagemaker.mlops.workflow import AthenaQueryConfig, AthenaStep, ParameterString | ||
|
|
||
| output_location = ParameterString(name="OutputLocation") |
There was a problem hiding this comment.
🟡 Low: The docstring example imports ParameterString from sagemaker.mlops.workflow, but that name is not exported from this package (it lives in sagemaker.core.workflow.parameters, and the __init__.py in this PR does not re-export it). A user copy-pasting this example would hit an ImportError. Note the tests correctly import it from sagemaker.core.workflow.parameters. Suggest updating the example import line accordingly.
rsareddy0329
left a comment
rsareddy0329
left a comment
There was a problem hiding this comment.
The serve unit test failures are not a regression as part of this PR.
pandas is not pinned in the SDK -- it's pulled in transitively (via mlflow) and only loosely constrained -- so the py311/py312 test envs now resolve the newly released pandas 3.1.0 today - https://pypi.org/project/pandas/3.1.0/ .
pandas 3.1.0 fails to import in those envs (its wheel is built against a newer NumPy ABI than the env resolves), so pytest errors at collection time, before any test runs. py310 is unaffected (it stays on 2.3.3, since pandas 3.x requires Python ≥ 3.11).
PR to temporarily unblock serve tests: #6524
Description
Adds AthenaStep, which runs an Amazon Athena query as a SageMaker Pipelines step. arguments is the Athena StartQueryExecution request as a dict and is emitted unchanged as the step's Arguments; the service validates it at CreatePipeline. properties follow Athena's GetQueryExecutionOutput shape, so downstream references use step.properties.QueryExecution..
sagemaker-core: load the botocore athena model into PropertiesMeta
sagemaker-mlops: new athena_step.py, StepTypeEnum.ATHENA, package export
Works with Run Conditions (StepDependency) through the base Step
Testing
Unit: 12 new tests in tests/unit/workflow/test_athena_step.py; full mlops workflow suite passes (497)
Integ: 3 tests in tests/integ/workflow/test_athena_step.py (successful query + S3 result, invalid SQL fails, full arguments round-trip) passed against a preprod stage