From 445096da2bd00857dd4f280e8d206e70c5be6899 Mon Sep 17 00:00:00 2001 From: Justus Perillieux Date: Thu, 16 Jul 2026 08:35:15 +0200 Subject: [PATCH] docs: add provisioning metadata examples --- python/examples/jobs.py | 16 ++++++++++ python/examples/jobs.rst | 7 +++++ python/examples/mmts.py | 25 ++++++++++++++++ python/examples/mmts.rst | 10 ++++++- python/examples/studios.py | 1 + python/examples/studios.rst | 4 +++ .../tests/core/test_provisioning_examples.py | 30 +++++++++++++++++++ 7 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 python/tests/core/test_provisioning_examples.py diff --git a/python/examples/jobs.py b/python/examples/jobs.py index 820610c5..a37b65b0 100755 --- a/python/examples/jobs.py +++ b/python/examples/jobs.py @@ -21,11 +21,19 @@ def _parser() -> argparse.ArgumentParser: studio_job.add_argument("--name", default="sdk-tutorial-job", help="Job name.") studio_job.add_argument("--machine", default="CPU", help="Machine type.") studio_job.add_argument("--command", default="python train.py --epochs 1", help="Command to run.") + studio_job.add_argument( + "--placement-group-id", + help="Existing placement group to join, for example from a Studio or another Job.", + ) image_job = subcommands.add_parser("image", help="Run a container-backed job.") image_job.add_argument("--name", default="sdk-image-job", help="Job name.") image_job.add_argument("--machine", default="CPU", help="Machine type.") image_job.add_argument("--image", default="python:3.11-slim", help="Docker image.") + image_job.add_argument( + "--placement-group-id", + help="Existing placement group to join, for example from a Studio or another Job.", + ) image_job.add_argument( "--command", default="python -c 'print(\"hello from a Lightning job\")'", @@ -49,12 +57,16 @@ def main() -> None: machine=Machine.from_str(args.machine), env={"RUN_MODE": "tutorial"}, command=args.command, + placement_group_id=args.placement_group_id, ) print(f"Job link: {job.link}") job.wait(interval=10, timeout=60 * 60, stop_on_timeout=True) print(job.json()) + print(f"Job resource id: {job.resource_id}") + print(f"Job private IP: {job.private_ip_address}") + print(f"Job placement group: {job.placement_group_id}") if job.status == Status.Failed: print(job.logs) # sdk-studio-job-end @@ -70,10 +82,14 @@ def main() -> None: command=args.command, env={"RUN_MODE": "tutorial"}, interruptible=True, + placement_group_id=args.placement_group_id, ) job.wait(interval=10) print(f"{job.name} finished with status {job.status}") + print(f"Job resource id: {job.resource_id}") + print(f"Job private IP: {job.private_ip_address}") + print(f"Job placement group: {job.placement_group_id}") # sdk-image-job-end diff --git a/python/examples/jobs.rst b/python/examples/jobs.rst index b14cebaa..b3ba14f2 100644 --- a/python/examples/jobs.rst +++ b/python/examples/jobs.rst @@ -59,6 +59,13 @@ Operational notes - ``Job.run`` creates a new job; ``Job("name", teamspace=...)`` fetches an existing one. +- Pass ``placement_group_id=...`` when a job must join an existing placement + group, for example to colocate with a running Studio, Job, or MMT. +- ``job.resource_id`` is the stable Lightning resource ID for the job. +- ``job.private_ip_address`` is populated when the job has a private worker + address available. +- ``job.placement_group_id`` reports the placement group associated with the + job, or ``None`` when the job is not tied to one. - ``job.logs`` is available after the job reaches a terminal state. - Studio-backed jobs must run in the same teamspace and cloud account as the Studio. diff --git a/python/examples/mmts.py b/python/examples/mmts.py index b039c164..f1a45c31 100755 --- a/python/examples/mmts.py +++ b/python/examples/mmts.py @@ -22,11 +22,19 @@ def _parser() -> argparse.ArgumentParser: studio_mmt.add_argument("--num-machines", type=int, default=2, help="Number of machines.") studio_mmt.add_argument("--machine", default="CPU", help="Machine type.") studio_mmt.add_argument("--command", default="python train_distributed.py --epochs 1", help="Command to run.") + studio_mmt.add_argument( + "--placement-group-id", + help="Existing placement group to join, for example from a Studio or Job.", + ) image_mmt = subcommands.add_parser("image", help="Run a container-backed MMT.") image_mmt.add_argument("--name", default="sdk-image-mmt", help="MMT name.") image_mmt.add_argument("--num-machines", type=int, default=2, help="Number of machines.") image_mmt.add_argument("--machine", default="L4", help="Machine type.") + image_mmt.add_argument( + "--placement-group-id", + help="Existing placement group to join, for example from a Studio or Job.", + ) image_mmt.add_argument( "--image", default="pytorch/pytorch:2.4.1-cuda12.1-cudnn9-runtime", @@ -56,15 +64,24 @@ def main() -> None: machine=Machine.from_str(args.machine), env={"RUN_MODE": "distributed"}, command=args.command, + placement_group_id=args.placement_group_id, ) print(f"MMT link: {mmt.link}") mmt.wait(interval=15, timeout=2 * 60 * 60, stop_on_timeout=True) print(mmt.json()) + print(f"MMT placement group: {mmt.placement_group_id}") if mmt.status == Status.Failed: for worker in mmt.machines: print(f"Worker {worker.name}: {worker.status}") + else: + for worker in mmt.machines: + print( + f"Worker rank={worker.rank} " + f"resource_id={worker.resource_id} " + f"private_ip={worker.private_ip_address}" + ) # sdk-studio-mmt-end elif args.example == "image": # sdk-image-mmt-start @@ -78,10 +95,18 @@ def main() -> None: machine=Machine.from_str(args.machine), command=args.command, interruptible=True, + placement_group_id=args.placement_group_id, ) mmt.wait(interval=15) print(f"{mmt.name} used {mmt.num_machines} machines") + print(f"MMT placement group: {mmt.placement_group_id}") + for worker in mmt.machines: + print( + f"Worker rank={worker.rank} " + f"resource_id={worker.resource_id} " + f"private_ip={worker.private_ip_address}" + ) # sdk-image-mmt-end diff --git a/python/examples/mmts.rst b/python/examples/mmts.rst index c4f4999b..31b72cca 100644 --- a/python/examples/mmts.rst +++ b/python/examples/mmts.rst @@ -54,7 +54,15 @@ Operational notes - ``MMT.run`` creates a new multi-machine job; ``MMT("name", teamspace=...)`` fetches an existing one. -- ``mmt.machines`` returns the per-machine job handles for detailed inspection. +- Pass ``placement_group_id=...`` when the whole MMT should join an existing + placement group, for example to colocate with a Studio or another workload. +- ``mmt.placement_group_id`` reports the placement group associated with the + multi-machine job, or ``None`` when the run is not tied to one. +- ``mmt.machines`` returns the per-machine job handles for detailed inspection, + sorted by ``job.rank``. +- Each member in ``mmt.machines`` exposes the same machine-level metadata as a + Job, including ``resource_id``, ``private_ip_address``, ``placement_group_id``, + and ``rank``. - ``MMT.run`` rejects ``num_machines`` values less than two. - Studio-backed MMT runs must use a Studio in the same teamspace and cloud account as the run. diff --git a/python/examples/studios.py b/python/examples/studios.py index 64554014..aadb28ed 100755 --- a/python/examples/studios.py +++ b/python/examples/studios.py @@ -34,6 +34,7 @@ def main() -> None: studio.start(machine=Machine.from_str(args.machine)) print(f"{studio.name} is {studio.status} on {studio.machine}") + print(f"Studio placement group: {studio.placement_group_id}") output = studio.run("python --version") print(output) diff --git a/python/examples/studios.rst b/python/examples/studios.rst index 7dd4a552..1b125485 100644 --- a/python/examples/studios.rst +++ b/python/examples/studios.rst @@ -60,6 +60,10 @@ Operational notes does not exist. Use ``create_ok=False`` when automation should fail instead of creating a new resource. - ``studio.start`` is blocking while compute is provisioned. +- ``studio.placement_group_id`` reports the active Studio compute placement + group, or ``None`` when placement metadata is not available. Use this value as + ``placement_group_id=...`` for Jobs or MMTs that must colocate with the + Studio. - ``studio.run`` requires the Studio to be running and raises if the command exits non-zero. - Jobs and MMTs can reuse a Studio environment, which is useful when the Studio diff --git a/python/tests/core/test_provisioning_examples.py b/python/tests/core/test_provisioning_examples.py new file mode 100644 index 00000000..f3a25766 --- /dev/null +++ b/python/tests/core/test_provisioning_examples.py @@ -0,0 +1,30 @@ +from pathlib import Path + +EXAMPLES_DIR = Path(__file__).resolve().parents[2] / "examples" + + +def test_job_example_documents_placement_metadata() -> None: + source = (EXAMPLES_DIR / "jobs.py").read_text() + + assert "--placement-group-id" in source + assert "placement_group_id=args.placement_group_id" in source + assert "job.resource_id" in source + assert "job.private_ip_address" in source + assert "job.placement_group_id" in source + + +def test_mmt_example_documents_ranked_worker_metadata() -> None: + source = (EXAMPLES_DIR / "mmts.py").read_text() + + assert "--placement-group-id" in source + assert "placement_group_id=args.placement_group_id" in source + assert "mmt.placement_group_id" in source + assert "worker.rank" in source + assert "worker.resource_id" in source + assert "worker.private_ip_address" in source + + +def test_studio_example_documents_placement_metadata() -> None: + source = (EXAMPLES_DIR / "studios.py").read_text() + + assert "studio.placement_group_id" in source