Document provisioning metadata examples

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
 
 

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.

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
 
 

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.

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)

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

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