feat: unify jobs + mmt under jobs command

python/lightning_sdk/api/job_api.py

@@ -143,6 +143,7 @@ def submit_job(
         reuse_snapshot: bool = True,
         scratch_disks: Optional[Dict[str, int]] = None,
         placement_group_id: Optional[str] = None,
+        num_machines: int = 1,
     ) -> V1Job:
         """Submit a v2 job and return the created job object.
 
@@ -164,10 +165,13 @@ def submit_job(
             reuse_snapshot: Whether to reuse the Studio's existing filesystem snapshot.
             scratch_disks: Optional mapping of scratch-disk mount paths to their sizes in GiB.
             placement_group_id: Optional placement group identifier for colocating the job.
+            num_machines: Must be 1 for single-machine jobs. Kept for parity with ``MMTApiV2.submit_job``.
 
         Returns:
             The newly created ``V1Job`` object.
         """
+        if num_machines != 1:
+            raise ValueError("JobApiV2 only supports single-machine jobs (num_machines=1)")
         if scratch_disks is not None:
             sanitized_scratch_disks = {}
             for k, v in scratch_disks.items():
@@ -342,7 +346,7 @@ def stop_job(self, job_id: str, teamspace_id: str) -> None:
                 break
             time.sleep(1)
 
-    def delete_job(self, job_id: str, teamspace_id: str, cloudspace_id: Optional[str]) -> None:
+    def delete_job(self, job_id: str, teamspace_id: str, cloudspace_id: Optional[str] = None) -> None:
         """Permanently delete a v2 job.
 
         Args:
@@ -622,3 +626,14 @@ def get_total_cost(self, job: V1Job) -> float:
             The total cost incurred by the job, expressed in US dollars.
         """
         return job.total_cost
+
+    def get_num_machines(self, job: V1Job) -> int:
+        """Return the number of machines for this job.
+
+        Single-machine jobs always report 1. Kept for parity with ``MMTApiV2.get_num_machines``.
+        """
+        return 1
+
+    def list_mmt_subjobs(self, job_id: str, teamspace_id: str) -> List[V1Job]:
+        """Single-machine jobs have no sub-jobs. Kept for parity with ``MMTApiV2.list_mmt_subjobs``."""
+        return []

python/lightning_sdk/api/mmt_api.py

@@ -51,6 +51,7 @@ def submit_job(
         max_runtime: Optional[int],
         reuse_snapshot: bool,
         placement_group_id: Optional[str] = None,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> V1MultiMachineJob:
         """Submit a v2 multi-machine job and return the created job object.
 
@@ -72,10 +73,13 @@ def submit_job(
             max_runtime: Maximum allowed runtime in seconds, or ``None`` for no limit.
             reuse_snapshot: Whether to reuse the Studio's existing filesystem snapshot.
             placement_group_id: Optional placement group identifier for colocating the job.
+            scratch_disks: Not supported for multi-machine jobs. Kept for parity with ``JobApiV2.submit_job``.
 
         Returns:
             The newly created ``V1MultiMachineJob`` object.
         """
+        if scratch_disks:
+            raise ValueError("scratch_disks are not supported for multi-machine jobs")
         body = self._create_mmt_body(
             name=name,
             num_machines=num_machines,
@@ -242,12 +246,13 @@ def stop_job(self, job_id: str, teamspace_id: str) -> None:
                 break
             time.sleep(1)
 
-    def delete_job(self, job_id: str, teamspace_id: str) -> None:
+    def delete_job(self, job_id: str, teamspace_id: str, cloudspace_id: Optional[str] = None) -> None:
         """Permanently delete a multi-machine job.
 
         Args:
             job_id: The unique identifier of the multi-machine job to delete.
             teamspace_id: The ID of the teamspace that owns the job.
+            cloudspace_id: Ignored. Kept for parity with ``JobApiV2.delete_job``.
         """
         self._client.jobs_service_delete_multi_machine_job(project_id=teamspace_id, id=job_id)
 

python/lightning_sdk/cli/job/inspect.py

@@ -6,7 +6,7 @@
 from rich.console import Console
 
 from lightning_sdk.cli.utils.logging import LightningCommand
-from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_teamspace
+from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace
 
 
 @click.command("inspect", cls=LightningCommand)
@@ -20,9 +20,19 @@
         "If not specified, uses the configured default teamspace."
     ),
 )
+@click.option("--rank", type=int, default=None, help="Inspect one machine in a multi-machine job.")
 @click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON (inspect always emits JSON).")
-def inspect_job(name: Optional[str] = None, teamspace: Optional[str] = None, as_json: bool = False) -> None:
+def inspect_job(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    rank: Optional[int] = None,
+    as_json: bool = False,
+) -> None:
     """Inspect a job for further details as JSON."""
     resolved_teamspace = resolve_teamspace(teamspace)
     job = resolve_job(name, resolved_teamspace)
+    if job.is_multi_machine is True and rank is not None:
+        job = resolve_job_machine(job, rank)
+    elif rank is not None:
+        raise click.UsageError("--rank is only supported for multi-machine jobs.")
     Console().print(job.json())

python/lightning_sdk/cli/job/list.py

@@ -1,10 +1,16 @@
 """Job list command."""
 
+from contextlib import suppress
 from typing import Optional
 
 import rich_click as click
+from rich.console import Console
+from rich.table import Table
 
+from lightning_sdk.cli.utils.json_output import echo_json
 from lightning_sdk.cli.utils.logging import LightningCommand
+from lightning_sdk.cli.utils.resource_resolution import resolve_teamspace
+from lightning_sdk.models import _list_teamspaces
 
 
 @click.command("list", cls=LightningCommand)
@@ -38,7 +44,55 @@ def list_jobs(
     sort_by: Optional[str] = None,
     as_json: bool = False,
 ) -> None:
-    """List jobs for a given teamspace."""
-    from lightning_sdk.cli.legacy.list import jobs
+    """List jobs for a given teamspace.
 
-    jobs.callback(teamspace=teamspace, all=all, sort_by=sort_by, as_json=as_json)
+    Includes both single- and multi-machine jobs.
+    """
+    resources = []
+    if all and not teamspace:
+        for teamspace_slug in _list_teamspaces():
+            resolved = resolve_teamspace(teamspace_slug)
+            resources.extend(resolved.jobs)
+    else:
+        resolved = resolve_teamspace(teamspace)
+        resources.extend(resolved.jobs)
+
+    rows = []
+    for job in resources:
+        job._prevent_refetch_latest = True
+        with suppress(RuntimeError):
+            rows.append(
+                {
+                    "name": job.name,
+                    "teamspace": f"{job.teamspace.owner.name}/{job.teamspace.name}",
+                    "studio": job.studio_name,
+                    "image": job.image,
+                    "status": str(job.status) if job.status is not None else None,
+                    "machine": str(job.machine),
+                    "num_machines": getattr(job, "num_machines", 1),
+                    "total_cost": round(job.total_cost, 3),
+                    "_cloud_account": str(getattr(job, "cloud_account", "") or ""),
+                }
+            )
+
+    sort_key = "_cloud_account" if sort_by == "cloud-account" else sort_by or "name"
+    rows.sort(key=lambda row: str(row.get(sort_key) or ""))
+    if as_json:
+        echo_json([{key: value for key, value in row.items() if not key.startswith("_")} for row in rows])
+        return
+
+    table = Table(pad_edge=True)
+    for column in ("Name", "Teamspace", "Studio", "Image", "Status", "Machine", "Num Machines", "Total Cost"):
+        table.add_column(column)
+    for row in rows:
+        table.add_row(
+            row["name"],
+            row["teamspace"],
+            row["studio"],
+            row["image"],
+            row["status"],
+            row["machine"],
+            str(row["num_machines"]),
+            f"{row['total_cost']:.3f}",
+        )
+    Console().print(table)

python/lightning_sdk/cli/job/logs.py

@@ -1,13 +1,14 @@
 """Job logs command."""
 
+from contextlib import suppress
 from typing import Optional
 
 import rich_click as click
 
 from lightning_sdk.api.logs_api import SEVERITIES
 from lightning_sdk.cli.utils.logging import LightningCommand
 from lightning_sdk.cli.utils.logs import LogSelection, read_logs, resolve_time
-from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_teamspace
+from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace
 
 
 @click.command("logs", cls=LightningCommand)
@@ -19,7 +20,7 @@
 )
 @click.option("--follow", "-f", is_flag=True, default=False, help="Stream new log lines as they are produced.")
 @click.option("--tail", type=int, default=None, help="Only show the last N lines.")
-@click.option("--rank", type=int, default=None, help="Distributed job rank to read from (running jobs only).")
+@click.option("--rank", type=int, default=None, help="Machine rank to read from in a multi-machine job.")
 @click.option("--timestamps", is_flag=True, default=False, help="Prepend each line with its ISO-8601 timestamp.")
 @click.option("--since", default=None, help='Only include lines at or after this time (e.g. "2h", RFC3339).')
 @click.option("--until", default=None, help='Only include lines at or before this time (e.g. "30m", RFC3339).')
@@ -49,15 +50,31 @@ def logs_job(
     Prints a snapshot of the logs available so far. Pass --follow to stream new
     lines from a running job until it finishes or you press Ctrl-C. --query and
     --severity are applied by the server, to both the snapshot and the stream.
+    Multi-machine logs are merged unless --rank selects one machine, which opens
+    that machine's per-job websocket (same path as a single-machine job with
+    --rank).
     """
     resolved_teamspace = resolve_teamspace(teamspace)
-    job = resolve_job(name, resolved_teamspace)
+    resource = resolve_job(name, resolved_teamspace)
+    selected_rank = resource.is_multi_machine is True and rank is not None
+    job = resolve_job_machine(resource, rank) if selected_rank else resource
 
     if as_json:
-        if rank is not None:
+        if rank is not None and not selected_rank:
             raise click.ClickException("--rank is not supported with --json.")
+        if job.is_multi_machine is True:
+            labels: dict = {}
+            with suppress(Exception):
+                labels = {machine.resource_id: machine.name for machine in job.machines}
+            selection = LogSelection(
+                teamspace_id=resolved_teamspace.id,
+                mmt_id=job.resource_id,
+                labels=labels,
+            )
+        else:
+            selection = LogSelection(teamspace_id=resolved_teamspace.id, job_ids=[job.resource_id])
         read_logs(
-            LogSelection(teamspace_id=resolved_teamspace.id, job_ids=[job.resource_id]),
+            selection,
             query=query,
             severity=severity,
             since=resolve_time(since, "--since"),
@@ -69,16 +86,20 @@ def logs_job(
         return
 
     try:
-        logs = job.logs(
-            follow=follow,
-            tail=tail,
-            rank=rank,
-            timestamps=timestamps,
-            since=resolve_time(since, "--since"),
-            until=resolve_time(until, "--until"),
-            query=query,
-            severity=severity,
-        )
+        log_options = {
+            "follow": follow,
+            "tail": tail,
+            "timestamps": timestamps,
+            "since": resolve_time(since, "--since"),
+            "until": resolve_time(until, "--until"),
+            "query": query,
+            "severity": severity,
+        }
+        if job.is_multi_machine is not True:
+            # Any non-None rank routes Job through the legacy per-job websocket (server-side
+            # tail). For a selected MMT machine the process rank on that node is 0.
+            log_options["rank"] = 0 if selected_rank else rank
+        logs = job.logs(**log_options)
         if follow:
             for line in logs:
                 click.echo(line)

python/lightning_sdk/cli/job/run.py

@@ -18,6 +18,14 @@
 
 @click.command("run", cls=LightningCommand)
 @click.option("--name", default=None, help="The name of the job. Needs to be unique within the teamspace.")
+@click.option(
+    "--num-machines",
+    "--num_machines",
+    default=1,
+    show_default=True,
+    type=click.IntRange(min=1),
+    help="The number of machines to run on.",
+)
 @click.option(
     "--machine",
     default="CPU",
@@ -119,6 +127,7 @@
 @click.option("--json", "as_json", is_flag=True, default=False, help="Output the created job as JSON.")
 def run_job(
     name: Optional[str] = None,
+    num_machines: int = 1,
     machine: str = "CPU",
     command: Optional[str] = None,
     studio: Optional[str] = None,
@@ -136,7 +145,10 @@ def run_job(
     path_mappings: str = "",
     as_json: bool = False,
 ) -> None:
-    """Run async workloads using a docker image or studio."""
+    """Run async workloads using a docker image or studio.
+
+    Pass --num-machines greater than 1 to run a multi-machine job.
+    """
     if not name:
         from datetime import datetime
 
@@ -159,23 +171,25 @@ def run_job(
     for value in env:
         env_dict.update(_resolve_envs(value))
 
-    job = Job.run(
-        name=name,
-        machine=machine_enum,
-        command=command,
-        studio=studio,
-        image=image,
-        teamspace=resolved_teamspace,
-        org=org,
-        user=user,
-        cloud=cloud,
-        env=env_dict,
-        interruptible=interruptible,
-        image_credentials=image_credentials,
-        cloud_account_auth=cloud_account_auth,
-        entrypoint=entrypoint,
-        path_mappings=path_mappings_dict,
-    )
+    run_kwargs = {
+        "name": name,
+        "machine": machine_enum,
+        "command": command,
+        "studio": studio,
+        "image": image,
+        "teamspace": resolved_teamspace,
+        "org": org,
+        "user": user,
+        "cloud": cloud,
+        "env": env_dict,
+        "interruptible": interruptible,
+        "image_credentials": image_credentials,
+        "cloud_account_auth": cloud_account_auth,
+        "entrypoint": entrypoint,
+        "path_mappings": path_mappings_dict,
+        "num_machines": num_machines,
+    }
+    job = Job.run(**run_kwargs)
 
     if as_json:
         echo_json(