Polyaxon v3 is coming →

Submit a Polyaxon workflow from AWS Lambda

Hand a trusted operation from a short AWS Lambda invocation to Polyaxon while accounting for retries and uncertain responses.

September 4, 2024by Polyaxon

Use AWS Lambda for a short handoff from an AWS event to a Polyaxon workload. For example, an application can emit a dataset-ready event, and the function can submit a predefined training or evaluation operation. The long-running work executes on Polyaxon; the function returns after submission.

This guide describes an application-owned Lambda function using the Polyaxon SDK. It does not deploy Lambda through Polyaxon or move training into Lambda. It is also separate from Lambda GPU Cloud, a different provider.

An AWS event invokes Lambda, which checks a durable submission ledger before handing a fixed operation to Polyaxon

Configure the function

  1. Package a compatible polyaxon client and boto3 with your function, together with a trusted operation.yaml. Follow AWS's Python handler guidance for packaging and handler configuration.
  2. Configure POLYAXON_HOST, POLYAXON_OWNER, and POLYAXON_PROJECT for the target deployment. These are application configuration, not values accepted from an untrusted event.
  3. Store a Polyaxon token in AWS Secrets Manager as a JSON object with a token field. Set POLYAXON_TOKEN_SECRET_ID to its secret identifier. Scope the function's execution role to that secret and any required encryption-key access.
  4. Ensure the function can reach both Secrets Manager and the Polyaxon API. For a private Polyaxon deployment, configure the appropriate VPC connectivity and DNS rather than exposing the API publicly for the function.
  5. Give the Polyaxon identity only the access required to submit the intended operation.

Submit a fixed operation

The following submission-only example expects an application event such as {"event_id": "dataset-ready-001"}. It is not an S3 event parser. It records a hashed event identifier for reconciliation but does not implement duplicate suppression; add the durable handoff described below before enabling a retrying production trigger. Package a real operation.yaml and pin the client versions used in your deployment.

import hashlib
import json
import os
from pathlib import Path

import boto3
from polyaxon.client import PolyaxonClient, RunClient


def lambda_handler(event, context):
    event_id = event.get("event_id")
    if not isinstance(event_id, str) or not event_id.strip() or len(event_id) > 256:
        raise ValueError("A nonempty event_id of at most 256 characters is required")

    secret = boto3.client("secretsmanager").get_secret_value(
        SecretId=os.environ["POLYAXON_TOKEN_SECRET_ID"]
    )
    token = json.loads(secret["SecretString"])["token"]
    client = RunClient(
        owner=os.environ["POLYAXON_OWNER"],
        project=os.environ["POLYAXON_PROJECT"],
        client=PolyaxonClient(token=token),
        manual_exceptions_handling=True,
    )
    event_tag = "event-" + hashlib.sha256(event_id.encode("utf-8")).hexdigest()
    run = client.create_from_polyaxonfile(
        polyaxonfile=str(Path(__file__).with_name("operation.yaml")),
        tags=["aws-lambda", event_tag],
    )
    return {"run_uuid": str(run.uuid)}

Set the handler to lambda_function.lambda_handler if you save the code in lambda_function.py. The packaged operation must already define its required inputs, image, and scheduling settings. Do not accept an arbitrary operation path or shell command from the event. The RunClient reference documents submission and subsequent status inspection.

Make retries safe

AWS event delivery can repeat. Before submission, use a durable event ledger with concurrency control; after acceptance, store the event-to-run UUID mapping. A function timeout can occur after Polyaxon creates the run but before the mapping is saved. Reconcile such uncertain submissions before retrying: a tag is a correlation aid, not an atomic idempotency key.

Keep failures visible to the trigger's retry or failure-handling policy, and monitor unresolved events. Use AWS's idempotency guidance, but do not assume a decorator or database write can make an external API call part of the same transaction.

Return after the handoff and observe training separately. For a multi-step workload, submit a Polyaxon pipeline rather than holding the Lambda invocation open while the run executes.