Skip to content

Agentic Orchestrator

"""Full agent: discover → submit → monitor → recover.

This is the "headline" example. Wire all the OSMO tools to a Strands Agent
and ask it to drive a workflow end-to-end.

Run:
    python examples/04_agentic_orchestrator.py "Submit ./train.yaml to an idle H100 pool"
"""
from __future__ import annotations

import sys

from strands import Agent

from strands_osmo import (
    osmo_cookbook_fetch,
    osmo_doctor,
    osmo_pool_list,
    osmo_resources_list,
    osmo_workflow_cancel,
    osmo_workflow_list,
    osmo_workflow_logs,
    osmo_workflow_render,
    osmo_workflow_status,
    osmo_workflow_submit,
    osmo_workflow_validate,
)


def main() -> None:
    query = (
        " ".join(sys.argv[1:])
        if len(sys.argv) > 1
        else (
            "Run osmo_doctor first. Then list available pools and tell me "
            "which would be best for a 4-GPU H100 training job. Don't submit "
            "anything yet."
        )
    )

    agent = Agent(tools=[
        osmo_doctor,
        osmo_pool_list,
        osmo_resources_list,
        osmo_cookbook_fetch,
        osmo_workflow_validate,
        osmo_workflow_render,
        osmo_workflow_submit,
        osmo_workflow_status,
        osmo_workflow_list,
        osmo_workflow_logs,
        osmo_workflow_cancel,
    ])

    print(f"💬 {query}\n")
    agent(query)


if __name__ == "__main__":
    main()