From a765c6d4e004c5e8990ba174334f678c8fa4a537 Mon Sep 17 00:00:00 2001 From: Teja Pulagam Date: Thu, 30 Jul 2026 21:51:17 +0000 Subject: [PATCH 1/6] feat: unify jobs + mmt under jobs command --- python/lightning_sdk/cli/job/delete.py | 4 +- python/lightning_sdk/cli/job/inspect.py | 21 ++++++- python/lightning_sdk/cli/job/list.py | 63 ++++++++++++++++++- python/lightning_sdk/cli/job/logs.py | 52 ++++++++++----- python/lightning_sdk/cli/job/run.py | 53 ++++++++++------ python/lightning_sdk/cli/job/ssh.py | 18 +++++- python/lightning_sdk/cli/job/stop.py | 4 +- .../cli/utils/resource_resolution.py | 40 +++++++++++- python/tests/cli/job/test_delete.py | 4 +- python/tests/cli/job/test_inspect.py | 25 +++++++- python/tests/cli/job/test_list.py | 44 +++++++++++++ python/tests/cli/job/test_logs.py | 55 ++++++++++++++-- python/tests/cli/job/test_run.py | 21 +++++++ python/tests/cli/job/test_ssh.py | 34 ++++++++-- python/tests/cli/job/test_stop.py | 2 +- .../cli/utils/test_resource_resolution.py | 27 ++++++++ 16 files changed, 406 insertions(+), 61 deletions(-) diff --git a/python/lightning_sdk/cli/job/delete.py b/python/lightning_sdk/cli/job/delete.py index 211eb316..afe6c28d 100644 --- a/python/lightning_sdk/cli/job/delete.py +++ b/python/lightning_sdk/cli/job/delete.py @@ -7,7 +7,7 @@ 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_job, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job_or_mmt, resolve_teamspace @click.command("delete", cls=LightningCommand) @@ -25,7 +25,7 @@ def delete_job(name: str, teamspace: Optional[str] = None, as_json: bool = False) -> None: """Delete a job.""" resolved_teamspace = resolve_teamspace(teamspace) - job = resolve_job(name, resolved_teamspace) + job = resolve_job_or_mmt(name, resolved_teamspace) job.delete() if as_json: echo_json({"name": job.name, "deleted": True}) diff --git a/python/lightning_sdk/cli/job/inspect.py b/python/lightning_sdk/cli/job/inspect.py index 856ca024..8e49eec6 100644 --- a/python/lightning_sdk/cli/job/inspect.py +++ b/python/lightning_sdk/cli/job/inspect.py @@ -6,7 +6,12 @@ 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_or_mmt, + resolve_mmt_machine, + resolve_teamspace, +) +from lightning_sdk.mmt import MMT @click.command("inspect", cls=LightningCommand) @@ -20,9 +25,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) + job = resolve_job_or_mmt(name, resolved_teamspace) + if isinstance(job, MMT) and rank is not None: + job = resolve_mmt_machine(job, rank) + elif rank is not None: + raise click.UsageError("--rank is only supported for multi-machine jobs.") Console().print(job.json()) diff --git a/python/lightning_sdk/cli/job/list.py b/python/lightning_sdk/cli/job/list.py index 29addc11..ca997fd9 100644 --- a/python/lightning_sdk/cli/job/list.py +++ b/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,58 @@ 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) + resources.extend(resolved.multi_machine_jobs) + else: + resolved = resolve_teamspace(teamspace) + resources.extend(resolved.jobs) + resources.extend(resolved.multi_machine_jobs) + + rows = [] + for job in resources: + job._prevent_refetch_latest = True + with suppress(RuntimeError): + studio = job.studio + rows.append( + { + "name": job.name, + "teamspace": f"{job.teamspace.owner.name}/{job.teamspace.name}", + "studio": studio.name if studio else None, + "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(job.cloud_account), + } + ) + + 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) diff --git a/python/lightning_sdk/cli/job/logs.py b/python/lightning_sdk/cli/job/logs.py index fd9e51d2..622d6134 100644 --- a/python/lightning_sdk/cli/job/logs.py +++ b/python/lightning_sdk/cli/job/logs.py @@ -1,5 +1,6 @@ """Job logs command.""" +from contextlib import suppress from typing import Optional import rich_click as click @@ -7,7 +8,12 @@ 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_or_mmt, + resolve_mmt_machine, + resolve_teamspace, +) +from lightning_sdk.mmt import MMT @click.command("logs", cls=LightningCommand) @@ -19,7 +25,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 +55,29 @@ 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. """ resolved_teamspace = resolve_teamspace(teamspace) - job = resolve_job(name, resolved_teamspace) + resource = resolve_job_or_mmt(name, resolved_teamspace) + selected_rank = isinstance(resource, MMT) and rank is not None + job = resolve_mmt_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 isinstance(job, MMT): + 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 +89,18 @@ 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 not isinstance(job, MMT): + log_options["rank"] = None if selected_rank else rank + logs = job.logs(**log_options) if follow: for line in logs: click.echo(line) diff --git a/python/lightning_sdk/cli/job/run.py b/python/lightning_sdk/cli/job/run.py index 8e69ff88..6ef3be83 100644 --- a/python/lightning_sdk/cli/job/run.py +++ b/python/lightning_sdk/cli/job/run.py @@ -10,6 +10,7 @@ from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace, teamspace_option from lightning_sdk.job import Job from lightning_sdk.machine import Machine +from lightning_sdk.mmt import MMT _MACHINE_VALUES = tuple( [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli] @@ -18,6 +19,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 +128,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 +146,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 +172,27 @@ 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, - ) + job_type = MMT if num_machines > 1 else Job + 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, + } + if job_type is MMT: + run_kwargs["num_machines"] = num_machines + job = job_type.run(**run_kwargs) if as_json: echo_json( diff --git a/python/lightning_sdk/cli/job/ssh.py b/python/lightning_sdk/cli/job/ssh.py index b3efe26e..56fdef27 100644 --- a/python/lightning_sdk/cli/job/ssh.py +++ b/python/lightning_sdk/cli/job/ssh.py @@ -6,8 +6,13 @@ import rich_click as click 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_or_mmt, + resolve_mmt_machine, + resolve_teamspace, +) from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal +from lightning_sdk.mmt import MMT from lightning_sdk.status import Status _SSH_HOST = "ssh.lightning.ai" @@ -31,27 +36,34 @@ def _ssh_user_for_job_id(job_id: str) -> str: "If not specified, uses the configured default teamspace." ), ) +@click.option("--rank", type=int, default=None, help="Machine rank for a multi-machine job. Defaults to 0.") def ssh_job( name: str, teamspace: Optional[str] = None, + rank: Optional[int] = None, ) -> None: """SSH into a running job. Example: lightning job ssh my-job """ - ssh_impl(name=name, teamspace=teamspace) + ssh_impl(name=name, teamspace=teamspace, rank=rank) def ssh_impl( name: Optional[str], teamspace: Optional[str], + rank: Optional[int] = None, ) -> None: if not name: raise click.UsageError("Missing job name. Pass NAME.") resolved_teamspace = resolve_teamspace(teamspace) - job = resolve_job(name, resolved_teamspace) + job = resolve_job_or_mmt(name, resolved_teamspace) + if isinstance(job, MMT): + job = resolve_mmt_machine(job, rank if rank is not None else 0) + elif rank is not None: + raise click.UsageError("--rank is only supported for multi-machine jobs.") if job.status != Status.Running: raise click.ClickException( diff --git a/python/lightning_sdk/cli/job/stop.py b/python/lightning_sdk/cli/job/stop.py index 3c4adcec..1632e9a5 100644 --- a/python/lightning_sdk/cli/job/stop.py +++ b/python/lightning_sdk/cli/job/stop.py @@ -7,7 +7,7 @@ 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_job, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job_or_mmt, resolve_teamspace @click.command("stop", cls=LightningCommand) @@ -25,7 +25,7 @@ def stop_job(name: str, teamspace: Optional[str] = None, as_json: bool = False) -> None: """Stop a job.""" resolved_teamspace = resolve_teamspace(teamspace) - job = resolve_job(name, resolved_teamspace) + job = resolve_job_or_mmt(name, resolved_teamspace) job.stop() if as_json: echo_json({"name": job.name, "status": "stopped"}) diff --git a/python/lightning_sdk/cli/utils/resource_resolution.py b/python/lightning_sdk/cli/utils/resource_resolution.py index 6ec19081..02b7642e 100644 --- a/python/lightning_sdk/cli/utils/resource_resolution.py +++ b/python/lightning_sdk/cli/utils/resource_resolution.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Optional, Union import rich_click as click @@ -65,6 +65,44 @@ def resolve_job(name: Optional[str], teamspace: Teamspace) -> Job: raise click.UsageError(f"Could not resolve job '{name}' in teamspace '{teamspace.name}'.") from ex +def resolve_job_or_mmt(name: Optional[str], teamspace: Teamspace) -> Union[Job, MMT]: + """Resolve a single- or multi-machine job by name.""" + if not name: + raise click.UsageError("Missing job name. Pass JOB.") + + try: + return Job(name=name, teamspace=teamspace) + except ValueError: + pass + + try: + return MMT(name=name, teamspace=teamspace) + except ValueError as ex: + raise click.UsageError(f"Could not resolve job '{name}' in teamspace '{teamspace.name}'.") from ex + + +def resolve_mmt_machine(mmt: MMT, rank: int) -> Job: + """Resolve one machine in a multi-machine job by rank.""" + machines = mmt.machines + if not machines: + raise click.ClickException(f"Job '{mmt.name}' has no machines.") + + expected = f"{mmt.name}-{rank}" + for machine in machines: + if machine.name == expected: + return machine + + prefix = f"{mmt.name}-" + available_ranks = [] + for machine in machines: + if machine.name.startswith(prefix): + suffix = machine.name[len(prefix) :] + if suffix.isdigit(): + available_ranks.append(int(suffix)) + available = ", ".join(str(value) for value in sorted(available_ranks)) + raise click.ClickException(f"Rank {rank} not found on job '{mmt.name}'. Available ranks: {available or 'none'}.") + + def resolve_mmt(name: Optional[str], teamspace: Teamspace) -> MMT: if not name: raise click.UsageError("Missing multi-machine job name. Pass JOB.") diff --git a/python/tests/cli/job/test_delete.py b/python/tests/cli/job/test_delete.py index 1af2bd8d..2fc1672c 100644 --- a/python/tests/cli/job/test_delete.py +++ b/python/tests/cli/job/test_delete.py @@ -48,7 +48,7 @@ def test_job_delete_resolves_exact_name() -> None: job = MagicMock() job.name = "train" with patch("lightning_sdk.cli.job.delete.resolve_teamspace", return_value=teamspace) as resolve_teamspace, patch( - "lightning_sdk.cli.job.delete.resolve_job", return_value=job + "lightning_sdk.cli.job.delete.resolve_job_or_mmt", return_value=job ) as resolve_job: result = CliRunner().invoke(delete_job, ["train", "--teamspace", "org/teamspace"]) @@ -67,7 +67,7 @@ def test_job_delete_json() -> None: job = MagicMock() job.name = "train" with patch("lightning_sdk.cli.job.delete.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.delete.resolve_job", return_value=job + "lightning_sdk.cli.job.delete.resolve_job_or_mmt", return_value=job ): result = CliRunner().invoke(delete_job, ["train", "--teamspace", "org/teamspace", "--json"]) diff --git a/python/tests/cli/job/test_inspect.py b/python/tests/cli/job/test_inspect.py index b8081967..2185cad0 100644 --- a/python/tests/cli/job/test_inspect.py +++ b/python/tests/cli/job/test_inspect.py @@ -2,6 +2,7 @@ from click.testing import CliRunner +from lightning_sdk.mmt import MMT from tests.cli.help import assert_help_contains, mock_command_logging @@ -54,7 +55,7 @@ def test_job_inspect_uses_positional_name() -> None: job.json.return_value = '{"name":"my-job"}' with patch("lightning_sdk.cli.job.inspect.resolve_teamspace", return_value=teamspace) as resolve_teamspace, patch( - "lightning_sdk.cli.job.inspect.resolve_job", return_value=job + "lightning_sdk.cli.job.inspect.resolve_job_or_mmt", return_value=job ) as resolve_job: result = CliRunner().invoke(inspect_job, ["my-job", "--teamspace", "org/teamspace"]) @@ -65,6 +66,28 @@ def test_job_inspect_uses_positional_name() -> None: job.json.assert_called_once_with() +@mock_command_logging +def test_job_inspect_selects_mmt_rank() -> None: + from lightning_sdk.cli.job.inspect import inspect_job + + teamspace = MagicMock() + mmt = MagicMock(spec=MMT) + machine = MagicMock() + machine.json.return_value = '{"name":"distributed-1"}' + with patch("lightning_sdk.cli.job.inspect.resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.cli.job.inspect.resolve_job_or_mmt", + return_value=mmt, + ), patch( + "lightning_sdk.cli.job.inspect.resolve_mmt_machine", + return_value=machine, + ) as resolve_rank: + result = CliRunner().invoke(inspect_job, ["distributed", "--rank", "1"]) + + assert result.exit_code == 0 + resolve_rank.assert_called_once_with(mmt, 1) + assert "distributed-1" in result.output + + @mock_command_logging def test_job_inspect_help_shows_positional_name() -> None: assert_help_contains("lightning job inspect --help", "Usage: lightning job inspect [OPTIONS] [NAME]") diff --git a/python/tests/cli/job/test_list.py b/python/tests/cli/job/test_list.py index d3df143e..aa59e410 100644 --- a/python/tests/cli/job/test_list.py +++ b/python/tests/cli/job/test_list.py @@ -1,5 +1,10 @@ +import json +from types import SimpleNamespace +from unittest.mock import patch + from click.testing import CliRunner +from lightning_sdk.cli.job.list import list_jobs from lightning_sdk.cli.legacy.list import jobs from tests.cli.help import assert_help_contains, mock_command_logging @@ -21,6 +26,45 @@ def test_jobs_list_help() -> None: assert_help_contains("lightning jobs list --help", "Usage: lightning jobs list", "List jobs for a given teamspace.") +@mock_command_logging +def test_job_list_includes_single_and_multi_machine_jobs() -> None: + owner = SimpleNamespace(name="org") + teamspace = SimpleNamespace(name="teamspace", owner=owner) + single = SimpleNamespace( + name="single", + teamspace=teamspace, + studio=None, + image="ubuntu", + status="Running", + machine="CPU", + total_cost=1.0, + cloud_account="default", + ) + multi = SimpleNamespace( + name="distributed", + teamspace=teamspace, + studio=None, + image="ubuntu", + status="Running", + machine="CPU", + num_machines=4, + total_cost=4.0, + cloud_account="default", + ) + teamspace.jobs = [single] + teamspace.multi_machine_jobs = [multi] + + with patch("lightning_sdk.cli.job.list.resolve_teamspace", return_value=teamspace): + result = CliRunner().invoke(list_jobs, ["--json"]) + + assert result.exit_code == 0, result.output + rows = json.loads(result.output) + assert [(row["name"], row["num_machines"]) for row in rows] == [ + ("distributed", 4), + ("single", 1), + ] + + @mock_command_logging def test_list_jobs_legacy_help() -> None: assert_help_contains( diff --git a/python/tests/cli/job/test_logs.py b/python/tests/cli/job/test_logs.py index 59c34ec3..7db4dad2 100644 --- a/python/tests/cli/job/test_logs.py +++ b/python/tests/cli/job/test_logs.py @@ -2,6 +2,7 @@ from click.testing import CliRunner +from lightning_sdk.mmt import MMT from tests.cli.help import assert_help_contains, mock_command_logging @@ -43,7 +44,7 @@ def test_job_logs_prints_snapshot() -> None: job = MagicMock() job.logs.return_value = "hello from the job\n42" with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=teamspace) as resolve_teamspace, patch( - "lightning_sdk.cli.job.logs.resolve_job", return_value=job + "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=job ) as resolve_job: result = CliRunner().invoke(logs_job, ["my-job", "--teamspace", "org/teamspace"]) @@ -65,7 +66,7 @@ def test_job_logs_follows_with_options() -> None: job = MagicMock() job.logs.return_value = iter(["line 1", "line 2"]) with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.logs.resolve_job", return_value=job + "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=job ): result = CliRunner().invoke( logs_job, @@ -79,6 +80,50 @@ def test_job_logs_follows_with_options() -> None: ) +@mock_command_logging +def test_job_logs_selects_mmt_rank() -> None: + from lightning_sdk.cli.job.logs import logs_job + + mmt = MagicMock(spec=MMT) + machine = MagicMock() + machine.logs.return_value = "rank one" + with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.job.logs.resolve_job_or_mmt", + return_value=mmt, + ), patch( + "lightning_sdk.cli.job.logs.resolve_mmt_machine", + return_value=machine, + ) as resolve_rank: + result = CliRunner().invoke(logs_job, ["distributed", "--rank", "1"]) + + assert result.exit_code == 0 + assert "rank one" in result.output + resolve_rank.assert_called_once_with(mmt, 1) + machine.logs.assert_called_once_with( + follow=False, tail=None, rank=None, timestamps=False, since=None, until=None, query=None, severity=None + ) + + +@mock_command_logging +def test_job_logs_merges_mmt_without_rank() -> None: + from lightning_sdk.cli.job.logs import logs_job + + mmt = MagicMock(spec=MMT) + mmt.logs.return_value = "[distributed-0] zero\n[distributed-1] one" + with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.job.logs.resolve_job_or_mmt", + return_value=mmt, + ): + result = CliRunner().invoke(logs_job, ["distributed"]) + + assert result.exit_code == 0 + assert "distributed-0" in result.output + assert "distributed-1" in result.output + mmt.logs.assert_called_once_with( + follow=False, tail=None, timestamps=False, since=None, until=None, query=None, severity=None + ) + + @mock_command_logging def test_job_logs_passes_filters() -> None: from lightning_sdk.cli.job.logs import logs_job @@ -87,7 +132,7 @@ def test_job_logs_passes_filters() -> None: job = MagicMock() job.logs.return_value = "boom" with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.logs.resolve_job", return_value=job + "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=job ): result = CliRunner().invoke(logs_job, ["my-job", "--query", "boom", "--severity", "error"]) @@ -101,7 +146,7 @@ def test_job_logs_passes_filters() -> None: def test_job_logs_rejects_unknown_severity() -> None: from lightning_sdk.cli.job.logs import logs_job - with patch("lightning_sdk.cli.job.logs.resolve_job") as resolve_job: + with patch("lightning_sdk.cli.job.logs.resolve_job_or_mmt") as resolve_job: result = CliRunner().invoke(logs_job, ["my-job", "--severity", "critical"]) assert result.exit_code != 0 @@ -116,7 +161,7 @@ def test_job_logs_reports_sdk_errors_cleanly() -> None: job = MagicMock() job.logs.side_effect = RuntimeError("Logs are not available while the job is Pending.") with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.logs.resolve_job", return_value=job + "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=job ): result = CliRunner().invoke(logs_job, ["my-job"]) diff --git a/python/tests/cli/job/test_run.py b/python/tests/cli/job/test_run.py index c1078a08..ad45b2ab 100644 --- a/python/tests/cli/job/test_run.py +++ b/python/tests/cli/job/test_run.py @@ -69,6 +69,27 @@ def test_run_job_with_cloud(monkeypatch): assert mock_job.run.call_args.kwargs["cloud"] == "aws" +@mock_command_logging +def test_run_job_with_multiple_machines_uses_mmt() -> None: + submitted = MagicMock(name="distributed") + submitted.name = "distributed" + with patch("lightning_sdk.cli.job.run.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.job.run.Job.run" + ) as run_single, patch( + "lightning_sdk.cli.job.run.MMT.run", + return_value=submitted, + ) as run_mmt: + result = CliRunner().invoke( + run_job, + ["--name", "distributed", "--image", "ubuntu", "--num-machines", "3"], + ) + + assert result.exit_code == 0, result.output + run_single.assert_not_called() + assert run_mmt.call_args.kwargs["num_machines"] == 3 + assert "Submitted job distributed." in result.output + + @pytest.mark.parametrize( ("input_mappings", "expected"), [ diff --git a/python/tests/cli/job/test_ssh.py b/python/tests/cli/job/test_ssh.py index 53610eb3..0820bb8f 100644 --- a/python/tests/cli/job/test_ssh.py +++ b/python/tests/cli/job/test_ssh.py @@ -5,6 +5,7 @@ from click.testing import CliRunner from lightning_sdk.cli.job.ssh import _ssh_user_for_job_id, ssh_impl, ssh_job +from lightning_sdk.mmt import MMT from lightning_sdk.status import Status from tests.cli.help import assert_help_contains, command_text, mock_command_logging @@ -27,7 +28,7 @@ def test_job_ssh_help() -> None: assert "Usage: lightning job ssh [OPTIONS] NAME" in result_text assert "SSH into a running job." in result_text assert "--teamspace" in result_text - assert "--rank" not in result_text + assert "--rank" in result_text @mock_command_logging @@ -42,7 +43,7 @@ def test_ssh_resolves_before_downloading_keys() -> None: "lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock(name="coding-model-training"), ), patch( - "lightning_sdk.cli.job.ssh.resolve_job", + "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", side_effect=click.UsageError("Could not resolve job 'missing'."), ), patch( "lightning_sdk.cli.job.ssh.configure_ssh_internal", @@ -60,7 +61,7 @@ def test_ssh_rejects_non_running_job() -> None: job.id = "job_01abc" with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch( - "lightning_sdk.cli.job.ssh.resolve_job", return_value=job + "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", return_value=job ), patch("lightning_sdk.cli.job.ssh.configure_ssh_internal") as configure, pytest.raises( click.ClickException, match="not Running" ): @@ -77,7 +78,7 @@ def test_ssh_runs_against_job_gateway_user() -> None: job.id = "job_01jj4hvvjj4zx1t1esm5az3zt7" with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.ssh.resolve_job", return_value=job + "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", return_value=job ) as resolve_job, patch( "lightning_sdk.cli.job.ssh.configure_ssh_internal", return_value="/tmp/lightning_rsa" ), patch("lightning_sdk.cli.job.ssh.subprocess.run") as run: @@ -88,6 +89,29 @@ def test_ssh_runs_against_job_gateway_user() -> None: run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_01jj4hvvjj4zx1t1esm5az3zt7@ssh.lightning.ai"]) +def test_ssh_selects_mmt_rank() -> None: + mmt = MagicMock(spec=MMT) + machine = MagicMock() + machine.name = "distributed-1" + machine.status = Status.Running + machine.id = "job_rank1" + + with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", + return_value=mmt, + ), patch( + "lightning_sdk.cli.job.ssh.resolve_mmt_machine", + return_value=machine, + ) as resolve_rank, patch( + "lightning_sdk.cli.job.ssh.configure_ssh_internal", + return_value="/tmp/lightning_rsa", + ), patch("lightning_sdk.cli.job.ssh.subprocess.run") as run: + ssh_impl(name="distributed", teamspace=None, rank=1) + + resolve_rank.assert_called_once_with(mmt, 1) + run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_rank1@ssh.lightning.ai"]) + + def test_ssh_retries_with_fresh_keys_on_failure() -> None: job = MagicMock() job.name = "train" @@ -98,7 +122,7 @@ def test_ssh_retries_with_fresh_keys_on_failure() -> None: run = MagicMock(side_effect=[OSError("ssh missing"), None]) with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch( - "lightning_sdk.cli.job.ssh.resolve_job", return_value=job + "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", return_value=job ), patch("lightning_sdk.cli.job.ssh.configure_ssh_internal", configure), patch( "lightning_sdk.cli.job.ssh.subprocess.run", run ): diff --git a/python/tests/cli/job/test_stop.py b/python/tests/cli/job/test_stop.py index 3bacf48a..9904bdf7 100644 --- a/python/tests/cli/job/test_stop.py +++ b/python/tests/cli/job/test_stop.py @@ -47,7 +47,7 @@ def test_job_stop_resolves_exact_name() -> None: job = MagicMock() job.name = "train" with patch("lightning_sdk.cli.job.stop.resolve_teamspace", return_value=teamspace) as resolve_teamspace, patch( - "lightning_sdk.cli.job.stop.resolve_job", return_value=job + "lightning_sdk.cli.job.stop.resolve_job_or_mmt", return_value=job ) as resolve_job: result = CliRunner().invoke(stop_job, ["train", "--teamspace", "org/teamspace"]) diff --git a/python/tests/cli/utils/test_resource_resolution.py b/python/tests/cli/utils/test_resource_resolution.py index a1db9b97..29337aac 100644 --- a/python/tests/cli/utils/test_resource_resolution.py +++ b/python/tests/cli/utils/test_resource_resolution.py @@ -8,7 +8,9 @@ join_teamspace_slug, resolve_cluster, resolve_job, + resolve_job_or_mmt, resolve_mmt, + resolve_mmt_machine, resolve_studio, resolve_teamspace, ) @@ -154,6 +156,31 @@ def test_resolve_job_converts_not_found_to_usage_error() -> None: resolve_job("train", MagicMock()) +def test_resolve_job_or_mmt_falls_back_to_mmt() -> None: + teamspace = MagicMock() + resolved = MagicMock() + with patch( + "lightning_sdk.cli.utils.resource_resolution.Job", + side_effect=ValueError("missing"), + ), patch( + "lightning_sdk.cli.utils.resource_resolution.MMT", + return_value=resolved, + ) as mmt: + assert resolve_job_or_mmt("distributed", teamspace) is resolved + mmt.assert_called_once_with(name="distributed", teamspace=teamspace) + + +def test_resolve_mmt_machine_uses_rank_suffix() -> None: + rank0 = MagicMock(name="rank-0") + rank0.name = "distributed-0" + rank1 = MagicMock(name="rank-1") + rank1.name = "distributed-1" + mmt = MagicMock(name="distributed", machines=(rank0, rank1)) + mmt.name = "distributed" + + assert resolve_mmt_machine(mmt, 1) is rank1 + + def test_resolve_mmt_requires_name() -> None: with pytest.raises(click.UsageError, match="JOB"): resolve_mmt(None, MagicMock()) From d77532ed4421094fb70a0bfb4052dd637b48fc2e Mon Sep 17 00:00:00 2001 From: Teja Pulagam Date: Fri, 31 Jul 2026 16:39:25 +0000 Subject: [PATCH 2/6] unify python classes + test fixes --- python/lightning_sdk/cli/job/inspect.py | 11 +- python/lightning_sdk/cli/job/list.py | 7 +- python/lightning_sdk/cli/job/logs.py | 23 +- python/lightning_sdk/cli/job/run.py | 7 +- python/lightning_sdk/cli/job/ssh.py | 11 +- .../cli/utils/resource_resolution.py | 34 +- python/lightning_sdk/job.py | 242 ++++++- python/lightning_sdk/mmt.py | 613 +----------------- python/lightning_sdk/pipeline/steps.py | 53 +- python/lightning_sdk/teamspace.py | 6 +- python/tests/cli/job/test_inspect.py | 3 +- python/tests/cli/job/test_list.py | 29 +- python/tests/cli/job/test_logs.py | 6 +- python/tests/cli/job/test_run.py | 11 +- python/tests/cli/job/test_ssh.py | 3 +- .../cli/utils/test_resource_resolution.py | 29 +- python/tests/core/test_job.py | 12 +- .../core/test_permissions_integration.py | 2 +- python/tests/core/test_pipeline.py | 9 + python/tests/core/test_teamspace.py | 11 +- python/tests/core/test_unified_job.py | 129 ++++ 21 files changed, 529 insertions(+), 722 deletions(-) create mode 100644 python/tests/core/test_unified_job.py diff --git a/python/lightning_sdk/cli/job/inspect.py b/python/lightning_sdk/cli/job/inspect.py index 8e49eec6..b329ad47 100644 --- a/python/lightning_sdk/cli/job/inspect.py +++ b/python/lightning_sdk/cli/job/inspect.py @@ -6,12 +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_or_mmt, - resolve_mmt_machine, - resolve_teamspace, -) -from lightning_sdk.mmt import MMT +from lightning_sdk.cli.utils.resource_resolution import resolve_job_machine, resolve_job_or_mmt, resolve_teamspace @click.command("inspect", cls=LightningCommand) @@ -36,8 +31,8 @@ def inspect_job( """Inspect a job for further details as JSON.""" resolved_teamspace = resolve_teamspace(teamspace) job = resolve_job_or_mmt(name, resolved_teamspace) - if isinstance(job, MMT) and rank is not None: - job = resolve_mmt_machine(job, rank) + 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()) diff --git a/python/lightning_sdk/cli/job/list.py b/python/lightning_sdk/cli/job/list.py index ca997fd9..3db95d5c 100644 --- a/python/lightning_sdk/cli/job/list.py +++ b/python/lightning_sdk/cli/job/list.py @@ -53,28 +53,25 @@ def list_jobs( for teamspace_slug in _list_teamspaces(): resolved = resolve_teamspace(teamspace_slug) resources.extend(resolved.jobs) - resources.extend(resolved.multi_machine_jobs) else: resolved = resolve_teamspace(teamspace) resources.extend(resolved.jobs) - resources.extend(resolved.multi_machine_jobs) rows = [] for job in resources: job._prevent_refetch_latest = True with suppress(RuntimeError): - studio = job.studio rows.append( { "name": job.name, "teamspace": f"{job.teamspace.owner.name}/{job.teamspace.name}", - "studio": studio.name if studio else None, + "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(job.cloud_account), + "_cloud_account": str(getattr(job, "cloud_account", "") or ""), } ) diff --git a/python/lightning_sdk/cli/job/logs.py b/python/lightning_sdk/cli/job/logs.py index 622d6134..b7f4481e 100644 --- a/python/lightning_sdk/cli/job/logs.py +++ b/python/lightning_sdk/cli/job/logs.py @@ -8,12 +8,7 @@ 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_or_mmt, - resolve_mmt_machine, - resolve_teamspace, -) -from lightning_sdk.mmt import MMT +from lightning_sdk.cli.utils.resource_resolution import resolve_job_machine, resolve_job_or_mmt, resolve_teamspace @click.command("logs", cls=LightningCommand) @@ -55,17 +50,19 @@ 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. + 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) resource = resolve_job_or_mmt(name, resolved_teamspace) - selected_rank = isinstance(resource, MMT) and rank is not None - job = resolve_mmt_machine(resource, rank) if selected_rank else resource + 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 and not selected_rank: raise click.ClickException("--rank is not supported with --json.") - if isinstance(job, MMT): + if job.is_multi_machine is True: labels: dict = {} with suppress(Exception): labels = {machine.resource_id: machine.name for machine in job.machines} @@ -98,8 +95,10 @@ def logs_job( "query": query, "severity": severity, } - if not isinstance(job, MMT): - log_options["rank"] = None if selected_rank else rank + 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: diff --git a/python/lightning_sdk/cli/job/run.py b/python/lightning_sdk/cli/job/run.py index 6ef3be83..1d86a7aa 100644 --- a/python/lightning_sdk/cli/job/run.py +++ b/python/lightning_sdk/cli/job/run.py @@ -10,7 +10,6 @@ from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace, teamspace_option from lightning_sdk.job import Job from lightning_sdk.machine import Machine -from lightning_sdk.mmt import MMT _MACHINE_VALUES = tuple( [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli] @@ -172,7 +171,6 @@ def run_job( for value in env: env_dict.update(_resolve_envs(value)) - job_type = MMT if num_machines > 1 else Job run_kwargs = { "name": name, "machine": machine_enum, @@ -189,10 +187,9 @@ def run_job( "cloud_account_auth": cloud_account_auth, "entrypoint": entrypoint, "path_mappings": path_mappings_dict, + "num_machines": num_machines, } - if job_type is MMT: - run_kwargs["num_machines"] = num_machines - job = job_type.run(**run_kwargs) + job = Job.run(**run_kwargs) if as_json: echo_json( diff --git a/python/lightning_sdk/cli/job/ssh.py b/python/lightning_sdk/cli/job/ssh.py index 56fdef27..768576fc 100644 --- a/python/lightning_sdk/cli/job/ssh.py +++ b/python/lightning_sdk/cli/job/ssh.py @@ -6,13 +6,8 @@ import rich_click as click from lightning_sdk.cli.utils.logging import LightningCommand -from lightning_sdk.cli.utils.resource_resolution import ( - resolve_job_or_mmt, - resolve_mmt_machine, - resolve_teamspace, -) +from lightning_sdk.cli.utils.resource_resolution import resolve_job_machine, resolve_job_or_mmt, resolve_teamspace from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal -from lightning_sdk.mmt import MMT from lightning_sdk.status import Status _SSH_HOST = "ssh.lightning.ai" @@ -60,8 +55,8 @@ def ssh_impl( resolved_teamspace = resolve_teamspace(teamspace) job = resolve_job_or_mmt(name, resolved_teamspace) - if isinstance(job, MMT): - job = resolve_mmt_machine(job, rank if rank is not None else 0) + if job.is_multi_machine is True: + job = resolve_job_machine(job, rank if rank is not None else 0) elif rank is not None: raise click.UsageError("--rank is only supported for multi-machine jobs.") diff --git a/python/lightning_sdk/cli/utils/resource_resolution.py b/python/lightning_sdk/cli/utils/resource_resolution.py index 02b7642e..d3487d40 100644 --- a/python/lightning_sdk/cli/utils/resource_resolution.py +++ b/python/lightning_sdk/cli/utils/resource_resolution.py @@ -1,4 +1,4 @@ -from typing import Optional, Union +from typing import Optional import rich_click as click @@ -65,34 +65,23 @@ def resolve_job(name: Optional[str], teamspace: Teamspace) -> Job: raise click.UsageError(f"Could not resolve job '{name}' in teamspace '{teamspace.name}'.") from ex -def resolve_job_or_mmt(name: Optional[str], teamspace: Teamspace) -> Union[Job, MMT]: +def resolve_job_or_mmt(name: Optional[str], teamspace: Teamspace) -> Job: """Resolve a single- or multi-machine job by name.""" - if not name: - raise click.UsageError("Missing job name. Pass JOB.") - - try: - return Job(name=name, teamspace=teamspace) - except ValueError: - pass - - try: - return MMT(name=name, teamspace=teamspace) - except ValueError as ex: - raise click.UsageError(f"Could not resolve job '{name}' in teamspace '{teamspace.name}'.") from ex + return resolve_job(name, teamspace) -def resolve_mmt_machine(mmt: MMT, rank: int) -> Job: +def resolve_job_machine(job: Job, rank: int) -> Job: """Resolve one machine in a multi-machine job by rank.""" - machines = mmt.machines + machines = job.machines if not machines: - raise click.ClickException(f"Job '{mmt.name}' has no machines.") + raise click.ClickException(f"Job '{job.name}' has no machines.") - expected = f"{mmt.name}-{rank}" + expected = f"{job.name}-{rank}" for machine in machines: if machine.name == expected: return machine - prefix = f"{mmt.name}-" + prefix = f"{job.name}-" available_ranks = [] for machine in machines: if machine.name.startswith(prefix): @@ -100,7 +89,12 @@ def resolve_mmt_machine(mmt: MMT, rank: int) -> Job: if suffix.isdigit(): available_ranks.append(int(suffix)) available = ", ".join(str(value) for value in sorted(available_ranks)) - raise click.ClickException(f"Rank {rank} not found on job '{mmt.name}'. Available ranks: {available or 'none'}.") + raise click.ClickException(f"Rank {rank} not found on job '{job.name}'. Available ranks: {available or 'none'}.") + + +def resolve_mmt_machine(mmt: MMT, rank: int) -> Job: + """Compatibility alias for resolving a machine by rank.""" + return resolve_job_machine(mmt, rank) def resolve_mmt(name: Optional[str], teamspace: Teamspace) -> MMT: diff --git a/python/lightning_sdk/job.py b/python/lightning_sdk/job.py index 65f9de35..109ba578 100644 --- a/python/lightning_sdk/job.py +++ b/python/lightning_sdk/job.py @@ -1,10 +1,11 @@ import warnings from pathlib import PurePath -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, TypedDict, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Tuple, TypedDict, Union from lightning_sdk.api.cloud_account_api import CloudAccountApi from lightning_sdk.api.job_api import JobApiV2 from lightning_sdk.api.logs_api import LogsApi +from lightning_sdk.api.mmt_api import MMTApiV2 from lightning_sdk.api.utils import AccessibleResource, _get_cloud_url, raise_access_error_if_not_allowed from lightning_sdk.status import Status from lightning_sdk.utils.logging import TrackCallsMeta @@ -121,7 +122,7 @@ class JobDict(TypedDict): class Job(metaclass=TrackCallsMeta): - """Submit and manage single-machine jobs on the Lightning AI Platform.""" + """Submit and manage jobs on the Lightning AI Platform.""" def __init__( self, @@ -131,6 +132,7 @@ def __init__( user: Union[str, "User", None] = None, *, _fetch_job: bool = True, + _resource_kind: Optional[str] = None, ) -> None: """Fetch already existing jobs. @@ -162,7 +164,12 @@ def __init__( self._job = None self._prevent_refetch_latest = False self._cloud_account_api = CloudAccountApi() - self._job_api = JobApiV2() + self._standalone_job_api = JobApiV2() + self._mmt_job_api: Optional[MMTApiV2] = None + self._resource_kind = _resource_kind or "standalone" + self._job_api: Union[JobApiV2, MMTApiV2] = self._standalone_job_api + if self._resource_kind == "multi": + self._set_resource_kind("multi") self._logs_api = LogsApi() if _fetch_job: @@ -175,6 +182,17 @@ def __init__( raise ValueError(f"Job {name} does not exist in Teamspace {teamspace.name}") from None raise + def _set_resource_kind(self, kind: str) -> None: + if kind not in ("standalone", "multi"): + raise ValueError(f"Unknown job resource kind: {kind}") + self._resource_kind = kind + if kind == "multi": + if self._mmt_job_api is None: + self._mmt_job_api = MMTApiV2() + self._job_api = self._mmt_job_api + else: + self._job_api = self._standalone_job_api + @classmethod def run( cls, @@ -197,12 +215,14 @@ def run( reuse_snapshot: bool = True, scratch_disks: Optional[Dict[str, int]] = None, placement_group_id: Optional[str] = None, + num_machines: int = 1, ) -> "Job": """Run async workloads using a docker image or a compute environment from your studio. Args: name: The name of the job. Needs to be unique within the teamspace. machine: The machine type to run the job on. + num_machines: The number of machines to run on. Defaults to one. command: The command to run inside your job. Required if using a studio. Optional if using an image. If not provided for images, will run the container entrypoint and default command. studio: The studio env to run the job with. Mutually exclusive with image. @@ -248,6 +268,10 @@ def run( if not name: raise ValueError("A job needs to have a name!") + if num_machines < 1: + raise ValueError("A job needs to run on at least one machine") + if num_machines > 1 and scratch_disks: + raise ValueError("scratch_disks are not supported for multi-machine jobs") if image is None: if not isinstance(studio, Studio): @@ -308,9 +332,11 @@ def run( entrypoint = None job = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False) + job._set_resource_kind("multi" if num_machines > 1 else "standalone") submit_cloud = cloud if cloud_account is None else None job._submit( + num_machines=num_machines, machine=machine, cloud=submit_cloud, command=command, @@ -350,7 +376,15 @@ def _submit( reuse_snapshot: bool = True, scratch_disks: Optional[Dict[str, int]] = None, placement_group_id: Optional[str] = None, + num_machines: int = 1, ) -> "Job": + if num_machines < 1: + raise ValueError("A job needs to run on at least one machine") + if self.is_multi_machine and num_machines <= 1: + raise ValueError("Multi-machine jobs need to run on at least two machines") + if self.is_multi_machine and scratch_disks: + raise ValueError("scratch_disks are not supported for multi-machine jobs") + if studio is not None: studio_id = studio._studio.id if image is not None: @@ -391,26 +425,29 @@ def _submit( if ".." in path.parts: raise ValueError("scratch_disk path cannot contain '..'") - submitted = self._job_api.submit_job( - name=self.name, - command=command, - cloud_account=cloud_account, - teamspace_id=self._teamspace.id, - studio_id=studio_id, - image=image, - machine=machine, - interruptible=interruptible, - env=env, - image_credentials=image_credentials, - cloud_account_auth=cloud_account_auth, - entrypoint=entrypoint, - path_mappings=path_mappings, - max_runtime=max_runtime, - reuse_snapshot=reuse_snapshot, - scratch_disks=scratch_disks, - placement_group_id=placement_group_id, - ) - if submitted.name != self._name: + submit_kwargs = { + "name": self.name, + "command": command, + "cloud_account": cloud_account, + "teamspace_id": self._teamspace.id, + "studio_id": studio_id, + "image": image, + "machine": machine, + "interruptible": interruptible, + "env": env, + "image_credentials": image_credentials, + "cloud_account_auth": cloud_account_auth, + "entrypoint": entrypoint, + "path_mappings": path_mappings, + "max_runtime": max_runtime, + "reuse_snapshot": reuse_snapshot, + "placement_group_id": placement_group_id, + } + if self.is_multi_machine: + submitted = self._job_api.submit_job(num_machines=num_machines, **submit_kwargs) + else: + submitted = self._job_api.submit_job(scratch_disks=scratch_disks, **submit_kwargs) + if not self.is_multi_machine and submitted.name != self._name: warnings.warn( f"Job name '{self._name}' was already taken in this teamspace; " f"the job was created as '{submitted.name}' instead.", @@ -428,11 +465,14 @@ def stop(self) -> None: self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id) def delete(self) -> None: - self._job_api.delete_job( - job_id=self._guaranteed_job.id, - teamspace_id=self._teamspace.id, - cloudspace_id=self._guaranteed_job.spec.cloudspace_id, - ) + if self.is_multi_machine: + self._job_api.delete_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id) + else: + self._job_api.delete_job( + job_id=self._guaranteed_job.id, + teamspace_id=self._teamspace.id, + cloudspace_id=self._guaranteed_job.spec.cloudspace_id, + ) def wait(self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False) -> None: import time @@ -485,6 +525,8 @@ def machine(self) -> Union["Machine", str]: @property def public_ip(self) -> Optional[str]: + if self.is_multi_machine: + return None try: return self._job.public_ip_address except AttributeError: @@ -501,6 +543,8 @@ def resource_id(self) -> Optional[str]: @property def private_ip_address(self) -> Optional[str]: + if self.is_multi_machine: + return None return self._guaranteed_job.private_ip_address @property @@ -509,10 +553,43 @@ def placement_group_id(self) -> Optional[str]: @property def rank(self) -> Optional[int]: + if self.is_multi_machine: + return None return self._guaranteed_job.spec.rank + @property + def is_multi_machine(self) -> bool: + """Whether this object represents a multi-machine parent job.""" + return self._resource_kind == "multi" + + @property + def num_machines(self) -> int: + """The number of machines allocated to this job.""" + if not self.is_multi_machine: + return 1 + return self._job_api.get_num_machines(self._guaranteed_job) + + @property + def machines(self) -> Tuple["Job", ...]: + """The rank-ordered machines in this job.""" + if not self.is_multi_machine: + return (self,) + + subjobs = sorted( + self._job_api.list_mmt_subjobs(self._guaranteed_job.id, self.teamspace.id), + key=lambda job: job.spec.rank, + ) + machines = [] + for subjob in subjobs: + job = Job(name=subjob.name, teamspace=self.teamspace, _fetch_job=False) + job._job = subjob + machines.append(job) + return tuple(machines) + @property def artifact_path(self) -> Optional[str]: + if self.is_multi_machine: + raise NotImplementedError if self._guaranteed_job.spec.image != "": if self._guaranteed_job.spec.artifacts_destination: ( @@ -527,12 +604,16 @@ def artifact_path(self) -> Optional[str]: @property def snapshot_path(self) -> Optional[str]: + if self.is_multi_machine: + raise NotImplementedError if self._guaranteed_job.spec.image != "": return None return f"/teamspace/jobs/{self._guaranteed_job.name}/snapshot" @property def share_path(self) -> Optional[str]: + if self.is_multi_machine: + return None raise NotImplementedError("Not implemented yet") @property @@ -579,6 +660,18 @@ def _compute_logs( severity: Optional[str] = None, ) -> Union[str, Iterator[str]]: """Fetch the logs, dispatching on job state. See :attr:`logs` for the public API.""" + if self.is_multi_machine: + return self._compute_multi_machine_logs( + follow=follow, + tail=tail, + rank=rank, + timestamps=timestamps, + since=since, + until=until, + query=query, + severity=severity, + ) + status = self.status if rank is not None: @@ -609,6 +702,72 @@ def _compute_logs( # keeping the return type consistent with the live-follow path. return iter(lines) if follow else "\n".join(lines) + def _compute_multi_machine_logs( + self, + *, + follow: bool, + tail: Optional[int], + rank: Optional[int], + timestamps: bool, + since: Optional[str], + until: Optional[str], + query: Optional[str], + severity: Optional[str], + ) -> Union[str, Iterator[str]]: + if rank is not None: + raise ValueError( + "`rank` is not supported on a multi-machine parent; " + "read a single machine with `job.machines[rank].logs` " + "(or `mmt.machines[rank].logs` through the compatibility API)." + ) + + status = self.status + if status not in (Status.Running, Status.Failed, Status.Completed, Status.Stopped): + raise RuntimeError(f"Logs are not available while the job is {status}.") + + lines = self._stream_multi_machine_entries( + follow=follow and status == Status.Running, + tail=tail, + timestamps=timestamps, + since=since, + until=until, + query=query, + severity=severity, + ) + if follow and status == Status.Running: + return lines + collected = list(lines) + return iter(collected) if follow else "\n".join(collected) + + def _stream_multi_machine_entries( + self, + *, + follow: bool, + tail: Optional[int], + timestamps: bool, + since: Optional[str], + until: Optional[str], + query: Optional[str], + severity: Optional[str], + ) -> Iterator[str]: + names = {machine._guaranteed_job.id: machine.name for machine in self.machines} + entries = self._logs_api.stream( + self.teamspace.id, + mmt_id=self._guaranteed_job.id, + since=since, + until=until, + query=query, + severity=severity, + follow=follow, + tail=tail, + tail_anchor=getattr(self._guaranteed_job, "stopped_at", None), + idle_timeout=None if follow else _RUNNING_LOGS_IDLE_TIMEOUT, + fallback_to_live=not follow, + stop=lambda: self.status in (Status.Stopped, Status.Completed, Status.Failed), + ) + for entry in entries: + yield entry.format(timestamps=timestamps, prefix=names.get(entry.resource_id, entry.resource_id)) + def _compute_logs_ranked( self, *, @@ -706,6 +865,12 @@ def _stream_logs( @property def link(self) -> str: + if self.is_multi_machine: + return ( + f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/" + f"jobs/{self.name}?app_id=mmt" + ) + mmt_name = self._job_api.get_mmt_name(self._guaranteed_job) if self._job_api.get_image_name(self._guaranteed_job): @@ -728,6 +893,11 @@ def link(self) -> str: def image(self) -> Optional[str]: return self._job_api.get_image_name(self._guaranteed_job) + @property + def studio_name(self) -> Optional[str]: + """The name of the studio this job runs in, without instantiating the Studio.""" + return self._job_api.get_studio_name(self._guaranteed_job) + @property def studio(self) -> Optional["Studio"]: from lightning_sdk.studio import Studio @@ -743,7 +913,21 @@ def command(self) -> str: def _update_internal_job(self) -> None: if getattr(self, "_job", None) is None: - self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) + if self.is_multi_machine: + self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) + return + + from lightning_sdk.lightning_cloud.openapi.rest import ApiException + + try: + self._job = self._standalone_job_api.get_job_by_name( + name=self._name, teamspace_id=self._teamspace.id + ) + except ApiException as ex: + if ex.status != 404: + raise + self._set_resource_kind("multi") + self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) return self._job = self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id) diff --git a/python/lightning_sdk/mmt.py b/python/lightning_sdk/mmt.py index 8b3fede0..3155d901 100644 --- a/python/lightning_sdk/mmt.py +++ b/python/lightning_sdk/mmt.py @@ -1,120 +1,56 @@ -import warnings -from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Protocol, Tuple, TypedDict, Union +from typing import TYPE_CHECKING, Dict, Optional, Protocol, Union -from lightning_sdk.api.cloud_account_api import CloudAccountApi -from lightning_sdk.api.logs_api import LogsApi -from lightning_sdk.api.mmt_api import MMTApiV2 -from lightning_sdk.api.utils import AccessibleResource, _get_cloud_url, raise_access_error_if_not_allowed -from lightning_sdk.job import _RUNNING_LOGS_IDLE_TIMEOUT, _Logs +from lightning_sdk.job import Job, JobDict from lightning_sdk.status import Status -from lightning_sdk.utils.logging import TrackCallsMeta -from lightning_sdk.utils.resolve import ( - _get_org_id, - _resolve_default_cloud_account, - _resolve_teamspace, - _setup_logger, - in_studio, - skip_studio_setup, -) if TYPE_CHECKING: - from lightning_sdk.job import Job from lightning_sdk.machine import CloudProvider, Machine from lightning_sdk.organization import Organization from lightning_sdk.studio import Studio from lightning_sdk.teamspace import Teamspace from lightning_sdk.user import User -_logger = _setup_logger(__name__) - __all__ = ["MMT", "MMTMachine"] -class MachineDict(TypedDict): - name: str - status: Status - machine: Union["Machine", str] - - class MMTMachine(Protocol): - """A single machine in multi-machine training.""" + """A single machine in a multi-machine job.""" @property - def name(self) -> str: - """The name of the individual machine. Usually corresponds to the rank. - - Returns: - str: The name of this machine instance. - """ - ... + def name(self) -> str: ... @property - def machine(self) -> Union["Machine", str]: - """The actual machine type this node is running on. - - Returns: - Union[Machine, str]: The machine type of this node. - """ - ... + def machine(self) -> Union["Machine", str]: ... @property - def artifact_path(self) -> Optional[str]: - """The path to the artifacts of this job. - - Returns: - Optional[str]: The artifact path, or None if not available. - """ - ... + def artifact_path(self) -> Optional[str]: ... @property - def status(self) -> Status: - """The status of this job. - - Returns: - Status: The current status of this machine's job. - """ - ... + def status(self) -> Status: ... @property - def resource_id(self) -> Optional[str]: - """The stable resource identifier for this machine.""" - ... + def resource_id(self) -> Optional[str]: ... @property - def private_ip_address(self) -> Optional[str]: - """The private IP address for this machine, if assigned.""" - ... + def private_ip_address(self) -> Optional[str]: ... @property - def placement_group_id(self) -> Optional[str]: - """The placement group identifier for this machine, if assigned.""" - ... + def placement_group_id(self) -> Optional[str]: ... @property - def rank(self) -> Optional[int]: - """The stable rank for this machine inside the multi-machine job.""" - ... + def rank(self) -> Optional[int]: ... @property - def logs(self) -> str: - """The logs of the given machine. - - Returns: - str: The complete logs from this machine's execution. - """ - ... + def logs(self) -> str: ... - def dict(self) -> MachineDict: - """Dict representation of the given machine. + def dict(self) -> JobDict: ... - Returns: - MachineDict: A dictionary containing the machine's name, status, and machine type. - """ - ... +class MMT(Job): + """Compatibility interface for multi-machine jobs. -class MMT(metaclass=TrackCallsMeta): - """Submit and manage multi-machine jobs on the Lightning AI Platform.""" + Multi-machine functionality is implemented by :class:`lightning_sdk.job.Job`. + """ def __init__( self, @@ -125,41 +61,19 @@ def __init__( *, _fetch_job: bool = True, ) -> None: - """Fetch already existing multi-machine jobs. - - Args: - name: the name of the job. - teamspace: the teamspace the job is part of. - org: the name of the organization owning the ``teamspace`` in case it is owned by an org. - Deprecated — pass the owner as part of ``teamspace`` instead, e.g. ``teamspace="owner/teamspace"``. - user: the name of the user owning the ``teamspace`` in case it is owned directly by a user instead - of an org. Deprecated — pass the owner as part of ``teamspace`` instead, - e.g. ``teamspace="owner/teamspace"``. - - Raises: - ValueError: If the teamspace cannot be resolved from the provided arguments, or if the job is not found - when ``_fetch_job=True``. - PermissionError: If the user does not have access to jobs in the given teamspace. - """ - teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user) - if teamspace is None: - raise ValueError( - "Cannot resolve the teamspace from provided arguments." - f" Got teamspace={teamspace}, org={org}, user={user}." + try: + super().__init__( + name=name, + teamspace=teamspace, + org=org, + user=user, + _fetch_job=_fetch_job, + _resource_kind="multi", ) - - raise_access_error_if_not_allowed(AccessibleResource.Jobs, teamspace_id=teamspace.id) - - self._teamspace = teamspace - self._name = name - self._job = None - self._prevent_refetch_latest = False - self._cloud_account_api = CloudAccountApi() - self._job_api = MMTApiV2() - self._logs_api = LogsApi() - - if _fetch_job: - self._update_internal_job() + except ValueError as ex: + resolved_teamspace = getattr(self, "_teamspace", None) + teamspace_name = getattr(resolved_teamspace, "name", teamspace) + raise ValueError(f"Multi-machine job {name} does not exist in Teamspace {teamspace_name}") from ex @classmethod def run( @@ -184,135 +98,21 @@ def run( reuse_snapshot: bool = True, placement_group_id: Optional[str] = None, ) -> "MMT": - """Run async workloads using a docker image across multiple machines. - - Args: - name: The name of the job. Needs to be unique within the teamspace. - num_machines: The number of machines to run on. - machine: The machine type to run the job on. - command: The command to run inside your job. Required if using a studio. Optional if using an image. - If not provided for images, will run the container entrypoint and default command. - studio: The studio env to run the job with. Mutually exclusive with image. - image: The docker image to run the job with. Mutually exclusive with studio. - teamspace: The teamspace the job should be associated with. Defaults to the current teamspace. - Accepts a bare name or an ``owner/teamspace`` slug. - org: The organization owning the teamspace, if any. Defaults to the current organization. - Deprecated — pass the owner as part of ``teamspace`` instead, e.g. ``teamspace="owner/teamspace"``. - user: The user owning the teamspace, if any. Defaults to the current user. - Deprecated — pass the owner as part of ``teamspace`` instead, e.g. ``teamspace="owner/teamspace"``. - cloud: Cloud provider or cloud account to run the job on. - env: Environment variables to set inside the job. - interruptible: Whether the job should run on interruptible instances. Cheaper but can be preempted. - image_credentials: Credentials secret name used to pull a private image. - cloud_account_auth: Whether to authenticate with the cloud account to pull the image. - Required if the registry is part of a cloud provider, such as ECR. - entrypoint: The entrypoint of your docker container. Defaults to ``sh -c`` which - just runs the provided command in a standard shell if a command is provided. - If no command is provided, it will run the pre-defined entrypoint of the provided image. - To use the pre-defined entrypoint of the provided image with a specified command, - set this to an empty string. - Only applicable when submitting docker jobs. - path_mappings: Maps container paths to data-connection paths in the form - ``{"": ":"}`` or ``{"": ""}`` - for the root of a connection. Only applicable when submitting docker jobs. - max_runtime: Duration in seconds to allocate the machine. Required for some top-end GCP machines. - Defaults to 3 hours. - reuse_snapshot: Whether to reuse a Studio snapshot when multiple jobs for the same Studio are - submitted. Turning this off may result in longer startup times. Defaults to True. - placement_group_id: Optional placement group identifier for colocating the job. - - Returns: - MMT: The newly submitted multi-machine job instance. - - Raises: - ValueError: If required arguments are missing or mutually exclusive arguments are both provided. - RuntimeError: If image and studio are both provided. - """ - from lightning_sdk.lightning_cloud.openapi.rest import ApiException - from lightning_sdk.studio import Studio - - cloud_account = _resolve_default_cloud_account(None) - if cloud is not None: - cloud_account = None - if num_machines <= 1: raise ValueError("Multi-Machine training cannot be run with less than 2 Machines") - if not name: - raise ValueError("A job needs to have a name!") - - if image is None: - if not isinstance(studio, Studio): - with skip_studio_setup(): - studio = Studio( - name=studio, - teamspace=teamspace, - org=org, - user=user, - cloud=cloud, - create_ok=False, - ) - - if teamspace is None: - teamspace = studio.teamspace - else: - teamspace_name = teamspace if isinstance(teamspace, str) else teamspace.name - if studio.teamspace.name != teamspace_name: - raise ValueError( - "Studio teamspace does not match provided teamspace. " - "Can only run jobs with Studio envs in the teamspace of that Studio." - ) - - if cloud_account is None: - cloud_account = studio.cloud_account - - if cloud_account != studio.cloud_account: - raise ValueError( - "Studio cloud_account does not match provided cloud_account. " - "Can only run jobs with Studio envs in the same cloud_account." - ) - - if image_credentials is not None: - raise ValueError("image_credentials is only supported when using a custom image") - - if cloud_account_auth: - raise ValueError("cloud_account_auth is only supported when using a custom image") - - if entrypoint is not None: - raise ValueError("Specifying the entrypoint has no effect for jobs with Studio envs.") - - else: - if studio is not None: - raise RuntimeError( - "image and studio are mutually exclusive as both define the environment to run the job in" - ) - - if cloud_account is None and cloud is None and in_studio(): - try: - with skip_studio_setup(): - resolve_studio = Studio(teamspace=teamspace, user=user, org=org) - cloud_account = resolve_studio.cloud_account - except (ValueError, ApiException): - warnings.warn("Could not infer cloud account from studio. Using teamspace default.") - - if command is not None and entrypoint is None: - entrypoint = "sh -c" - elif entrypoint == "" or entrypoint is None: - entrypoint = None - - mmt = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False) - submit_cloud = cloud if cloud_account is None else None - - mmt._submit( - num_machines=num_machines, + return super().run( + name=name, machine=machine, - cloud=submit_cloud, + cloud=cloud, command=command, studio=studio, image=image, + teamspace=teamspace, + org=org, + user=user, env=env, interruptible=interruptible, - cloud_account=cloud_account, image_credentials=image_credentials, cloud_account_auth=cloud_account_auth, entrypoint=entrypoint, @@ -320,348 +120,5 @@ def run( max_runtime=max_runtime, reuse_snapshot=reuse_snapshot, placement_group_id=placement_group_id, - ) - - _logger.info(f"Multi-Machine Job was successfully launched. View it at {mmt.link}") - return mmt - - def _submit( - self, - num_machines: int, - machine: Union["Machine", str], - cloud: Optional[Union["CloudProvider", str]] = None, - command: Optional[str] = None, - studio: Optional["Studio"] = None, - image: Optional[str] = None, - env: Optional[Dict[str, str]] = None, - interruptible: bool = False, - cloud_account: Optional[str] = None, - image_credentials: Optional[str] = None, - cloud_account_auth: bool = False, - entrypoint: Optional[str] = None, - path_mappings: Optional[Dict[str, str]] = None, - max_runtime: Optional[int] = None, - reuse_snapshot: bool = True, - placement_group_id: Optional[str] = None, - ) -> "MMT": - if studio is not None: - studio_id = studio._studio.id - if image is not None: - raise ValueError( - "image and studio are mutually exclusive as both define the environment to run the job in" - ) - if command is None: - raise ValueError("command is required when using a studio") - else: - studio_id = None - if image is None: - raise ValueError("either image or studio must be provided") - - cloud_account = self._cloud_account_api.resolve_cloud_account( - self._teamspace.id, - cloud=cloud or cloud_account, - default_cloud_account=self._teamspace.default_cloud_account, - ) - - submitted = self._job_api.submit_job( - name=self.name, num_machines=num_machines, - command=command, - cloud_account=cloud_account, - teamspace_id=self._teamspace.id, - studio_id=studio_id, - image=image, - machine=machine, - interruptible=interruptible, - env=env, - image_credentials=image_credentials, - cloud_account_auth=cloud_account_auth, - entrypoint=entrypoint, - path_mappings=path_mappings, - max_runtime=max_runtime, - reuse_snapshot=reuse_snapshot, - placement_group_id=placement_group_id, ) - self._job = submitted - self._name = submitted.name - return self - - @property - def machines(self) -> Tuple["Job", ...]: - from lightning_sdk.job import Job - - subjobs = sorted( - self._job_api.list_mmt_subjobs(self._guaranteed_job.id, self.teamspace.id), - key=lambda job: job.spec.rank, - ) - machines = [] - for subjob in subjobs: - job = Job(name=subjob.name, teamspace=self.teamspace, _fetch_job=False) - job._job = subjob - machines.append(job) - return tuple(machines) - - def stop(self) -> None: - if self.status in (Status.Stopped, Status.Completed, Status.Failed): - return - self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id) - - def delete(self) -> None: - self._job_api.delete_job( - job_id=self._guaranteed_job.id, - teamspace_id=self._teamspace.id, - ) - - def wait(self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False) -> None: - import time - - start = time.time() - while True: - if self.status in (Status.Completed, Status.Stopped, Status.Failed): - break - - if timeout is not None and time.time() - start > timeout: - if stop_on_timeout: - self.stop() - raise TimeoutError("Job didn't finish within the provided timeout.") - - time.sleep(interval) - - async def async_wait( - self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False - ) -> None: - import asyncio - - start = asyncio.get_event_loop().time() - while True: - if self.status in (Status.Completed, Status.Stopped, Status.Failed): - break - - if timeout is not None and asyncio.get_event_loop().time() - start > timeout: - if stop_on_timeout: - self.stop() - raise TimeoutError("Job didn't finish within the provided timeout.") - - await asyncio.sleep(interval) - - @property - def status(self) -> Status: - return self._job_api._job_state_to_external(self._latest_job.state) - - @property - def id(self) -> Optional[str]: - """The multi-machine job's unique identifier.""" - return self._job.id if self._job is not None else None - - @property - def placement_group_id(self) -> Optional[str]: - return self._guaranteed_job.spec.placement_group_id - - @property - def artifact_path(self) -> Optional[str]: - raise NotImplementedError - - @property - def snapshot_path(self) -> Optional[str]: - raise NotImplementedError - - @property - def share_path(self) -> Optional[str]: - return None - - @property - def machine(self) -> Union["Machine", str]: - return self._job_api._get_job_machine_from_spec( - self._guaranteed_job.spec, - self.teamspace.id, - _get_org_id(self.teamspace), - ) - - def _update_internal_job(self) -> None: - if getattr(self, "_job", None) is None: - from lightning_sdk.lightning_cloud.openapi.rest import ApiException - - try: - self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) - except ApiException as ex: - if ex.status != 404: - raise - raise ValueError( - f"Multi-machine job {self._name} does not exist in Teamspace {self._teamspace.name}" - ) from ex - return - - self._job = self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id) - - @property - def name(self) -> str: - return self._name - - @property - def resource_id(self) -> Optional[str]: - return self._guaranteed_job.id - - @property - def teamspace(self) -> "Teamspace": - return self._teamspace - - @property - def link(self) -> str: - return f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=mmt" - - @property - def image(self) -> Optional[str]: - return self._job_api.get_image_name(self._guaranteed_job) - - @property - def studio(self) -> Optional["Studio"]: - from lightning_sdk.studio import Studio - - studio_name = self._job_api.get_studio_name(self._guaranteed_job) - if not studio_name: - return None - return Studio(studio_name, teamspace=self.teamspace) - - @property - def command(self) -> str: - return self._job_api.get_command(self._guaranteed_job) - - @property - def num_machines(self) -> int: - return self._job_api.get_num_machines(self._guaranteed_job) - - @property - def logs(self) -> _Logs: - """The logs of every machine, merged into one timeline. - - Use it as a value for a snapshot of the logs up to now:: - - print(mmt.logs) - - or call it to pass options and/or follow the logs live:: - - recent = mmt.logs(tail=100) # snapshot of the last 100 lines - for line in mmt.logs(follow=True): # stream new lines as they arrive - print(line) - - Options: - - - ``follow``: Keep the stream open and yield new lines as they are produced. - Returns an iterator of lines instead of a string. - - ``tail``: Only include the last N lines. - - ``timestamps``: Prepend each line with its ISO-8601 timestamp. - - ``since``/``until``: Only include lines within this RFC3339 time range. - - ``query``: Only include lines containing every whitespace-separated term. - - ``severity``: Only include lines at or above this level (``error``, ``warning``, - ``info`` or ``debug``). - - Every line is labelled with the machine it came from. To read a single machine, use - ``mmt.machines[rank].logs``. - """ - return _Logs(self._compute_logs) - - def _compute_logs( - self, - *, - follow: bool = False, - tail: Optional[int] = None, - rank: Optional[int] = None, - timestamps: bool = False, - since: Optional[str] = None, - until: Optional[str] = None, - query: Optional[str] = None, - severity: Optional[str] = None, - ) -> Union[str, Iterator[str]]: - """Fetch the merged logs of every machine. See :attr:`logs` for the public API.""" - if rank is not None: - raise ValueError("`rank` is not supported here; read a single machine with `mmt.machines[rank].logs`.") - - status = self.status - if status not in (Status.Running, Status.Failed, Status.Completed, Status.Stopped): - raise RuntimeError(f"Logs are not available while the job is {status}.") - - lines = self._stream_entries( - follow=follow and status == Status.Running, - tail=tail, - timestamps=timestamps, - since=since, - until=until, - query=query, - severity=severity, - ) - if follow and status == Status.Running: - return lines - collected = list(lines) - return iter(collected) if follow else "\n".join(collected) - - def _stream_entries( - self, - *, - follow: bool, - tail: Optional[int], - timestamps: bool, - since: Optional[str] = None, - until: Optional[str] = None, - query: Optional[str] = None, - severity: Optional[str] = None, - ) -> Iterator[str]: - """Yield formatted log lines for every machine, labelled with the machine they came from.""" - names = {machine._guaranteed_job.id: machine.name for machine in self.machines} - entries = self._logs_api.stream( - self.teamspace.id, - mmt_id=self._guaranteed_job.id, - since=since, - until=until, - query=query, - severity=severity, - follow=follow, - tail=tail, - # A finished job's last lines sit at its stop time, so start the tail search there - # instead of walking back from now through a job that ran days ago. - tail_anchor=getattr(self._guaranteed_job, "stopped_at", None), - idle_timeout=None if follow else _RUNNING_LOGS_IDLE_TIMEOUT, - # A running job whose logs are not in the current storage format yet has no saved - # history; tail its live stream so a snapshot still shows something. - fallback_to_live=not follow, - stop=lambda: self.status in (Status.Stopped, Status.Completed, Status.Failed), - ) - for entry in entries: - yield entry.format(timestamps=timestamps, prefix=names.get(entry.resource_id, entry.resource_id)) - - def dict(self) -> Dict[str, object]: - studio = self.studio - - return { - "name": self.name, - "teamspace": f"{self.teamspace.owner.name}/{self.teamspace.name}", - "studio": studio.name if studio else None, - "image": self.image, - "command": self.command, - "status": self.status, - "machine": self.machine, - "total_cost": self.total_cost, - } - - def json(self) -> str: - import json - - return json.dumps(self.dict(), indent=4, sort_keys=True, default=str) - - @property - def _guaranteed_job(self) -> Any: - if getattr(self, "_job", None) is None: - self._update_internal_job() - - return self._job - - @property - def total_cost(self) -> float: - return self._job_api.get_total_cost(self._latest_job) - - @property - def _latest_job(self) -> Any: - if self._prevent_refetch_latest: - return self._guaranteed_job - - self._update_internal_job() - return self._job diff --git a/python/lightning_sdk/pipeline/steps.py b/python/lightning_sdk/pipeline/steps.py index e41d3427..17fbf60f 100644 --- a/python/lightning_sdk/pipeline/steps.py +++ b/python/lightning_sdk/pipeline/steps.py @@ -207,6 +207,7 @@ def __init__( reuse_snapshot: bool = True, scratch_disks: Optional[Dict[str, int]] = None, placement_group_id: Optional[str] = None, + num_machines: int = 1, ) -> None: """Configure a job step in a pipeline. @@ -231,8 +232,11 @@ def __init__( reuse_snapshot: Whether to reuse a studio snapshot across jobs. Defaults to True. scratch_disks: Extra volumes to mount under ``/teamspace/scratch``. placement_group_id: Optional placement group identifier for colocating the job. + num_machines: Number of machines to allocate. Defaults to one. """ + if num_machines < 1: + raise ValueError("A job needs to run on at least one machine") self.name = name self.machine = machine or Machine.CPU self.command = command @@ -254,6 +258,7 @@ def __init__( self.reuse_snapshot = reuse_snapshot self.scratch_disks = scratch_disks self.placement_group_id = placement_group_id + self.num_machines = num_machines def to_proto( self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem] @@ -286,26 +291,36 @@ def to_proto( _validate_cloud_account(cloud_account, resolved_cloud_account, shared_filesystem) - body = JobApiV2._create_job_body( - name=self.name, - command=self.command, - cloud_account=resolved_cloud_account or cloud_account, - studio_id=studio._studio.id if isinstance(studio, Studio) else None, - image=self.image, - machine=self.machine, - interruptible=self.interruptible, - env=self.env, - image_credentials=self.image_credentials, - cloud_account_auth=self.cloud_account_auth, - entrypoint=self.entrypoint, - path_mappings=self.path_mappings, - max_runtime=self.max_runtime, - machine_image_version=machine_image_version, - reuse_snapshot=self.reuse_snapshot, - scratch_disks=self.scratch_disks, - placement_group_id=self.placement_group_id, - ) + body_kwargs = { + "name": self.name, + "command": self.command, + "cloud_account": resolved_cloud_account or cloud_account, + "studio_id": studio._studio.id if isinstance(studio, Studio) else None, + "image": self.image, + "machine": self.machine, + "interruptible": self.interruptible, + "env": self.env, + "image_credentials": self.image_credentials, + "cloud_account_auth": self.cloud_account_auth, + "entrypoint": self.entrypoint, + "path_mappings": self.path_mappings, + "max_runtime": self.max_runtime, + "machine_image_version": machine_image_version, + "reuse_snapshot": self.reuse_snapshot, + "placement_group_id": self.placement_group_id, + } + if self.num_machines > 1: + if self.scratch_disks: + raise ValueError("scratch_disks are not supported for multi-machine jobs") + body = MMTApiV2._create_mmt_body(num_machines=self.num_machines, **body_kwargs) + return V1PipelineStep( + name=self.name, + type=V1PipelineStepType.MMT, + wait_for=_to_wait_for(self.wait_for), + mmt=body, + ) + body = JobApiV2._create_job_body(scratch_disks=self.scratch_disks, **body_kwargs) return V1PipelineStep( name=self.name, type=V1PipelineStepType.JOB, diff --git a/python/lightning_sdk/teamspace.py b/python/lightning_sdk/teamspace.py index a531e750..d03d911b 100644 --- a/python/lightning_sdk/teamspace.py +++ b/python/lightning_sdk/teamspace.py @@ -278,7 +278,7 @@ def clusters(self) -> List[str]: @property def jobs(self) -> Tuple["Job", ...]: - """All single-machine jobs in this teamspace. + """All standalone and multi-machine jobs in this teamspace. Returns: tuple[Job, ...]: Every Job belonging to this teamspace. @@ -296,6 +296,10 @@ def jobs(self) -> Tuple["Job", ...]: job = Job(name=j2.name, teamspace=self, _fetch_job=False) job._job = j2 jobs.append(job) + for m2 in self._teamspace_api.list_mmts(teamspace_id=self.id): + job = Job(name=m2.name, teamspace=self, _fetch_job=False, _resource_kind="multi") + job._job = m2 + jobs.append(job) return tuple(jobs) diff --git a/python/tests/cli/job/test_inspect.py b/python/tests/cli/job/test_inspect.py index 2185cad0..af294994 100644 --- a/python/tests/cli/job/test_inspect.py +++ b/python/tests/cli/job/test_inspect.py @@ -72,13 +72,14 @@ def test_job_inspect_selects_mmt_rank() -> None: teamspace = MagicMock() mmt = MagicMock(spec=MMT) + mmt.is_multi_machine = True machine = MagicMock() machine.json.return_value = '{"name":"distributed-1"}' with patch("lightning_sdk.cli.job.inspect.resolve_teamspace", return_value=teamspace), patch( "lightning_sdk.cli.job.inspect.resolve_job_or_mmt", return_value=mmt, ), patch( - "lightning_sdk.cli.job.inspect.resolve_mmt_machine", + "lightning_sdk.cli.job.inspect.resolve_job_machine", return_value=machine, ) as resolve_rank: result = CliRunner().invoke(inspect_job, ["distributed", "--rank", "1"]) diff --git a/python/tests/cli/job/test_list.py b/python/tests/cli/job/test_list.py index aa59e410..fdd8b84a 100644 --- a/python/tests/cli/job/test_list.py +++ b/python/tests/cli/job/test_list.py @@ -26,33 +26,37 @@ def test_jobs_list_help() -> None: assert_help_contains("lightning jobs list --help", "Usage: lightning jobs list", "List jobs for a given teamspace.") -@mock_command_logging -def test_job_list_includes_single_and_multi_machine_jobs() -> None: +def _teamspace_with_jobs() -> SimpleNamespace: + """Build a teamspace whose jobs expose only the attributes real Job/MMT objects have.""" owner = SimpleNamespace(name="org") teamspace = SimpleNamespace(name="teamspace", owner=owner) single = SimpleNamespace( name="single", teamspace=teamspace, - studio=None, + studio_name=None, image="ubuntu", status="Running", machine="CPU", total_cost=1.0, - cloud_account="default", ) multi = SimpleNamespace( name="distributed", teamspace=teamspace, - studio=None, + studio_name=None, image="ubuntu", status="Running", machine="CPU", num_machines=4, total_cost=4.0, - cloud_account="default", ) - teamspace.jobs = [single] + teamspace.jobs = [single, multi] teamspace.multi_machine_jobs = [multi] + return teamspace + + +@mock_command_logging +def test_job_list_includes_single_and_multi_machine_jobs() -> None: + teamspace = _teamspace_with_jobs() with patch("lightning_sdk.cli.job.list.resolve_teamspace", return_value=teamspace): result = CliRunner().invoke(list_jobs, ["--json"]) @@ -65,6 +69,17 @@ def test_job_list_includes_single_and_multi_machine_jobs() -> None: ] +@mock_command_logging +def test_job_list_sort_by_cloud_account_without_attribute() -> None: + teamspace = _teamspace_with_jobs() + + with patch("lightning_sdk.cli.job.list.resolve_teamspace", return_value=teamspace): + result = CliRunner().invoke(list_jobs, ["--sort-by", "cloud-account", "--json"]) + + assert result.exit_code == 0, result.output + assert {row["name"] for row in json.loads(result.output)} == {"single", "distributed"} + + @mock_command_logging def test_list_jobs_legacy_help() -> None: assert_help_contains( diff --git a/python/tests/cli/job/test_logs.py b/python/tests/cli/job/test_logs.py index 7db4dad2..7839bd2d 100644 --- a/python/tests/cli/job/test_logs.py +++ b/python/tests/cli/job/test_logs.py @@ -85,13 +85,14 @@ def test_job_logs_selects_mmt_rank() -> None: from lightning_sdk.cli.job.logs import logs_job mmt = MagicMock(spec=MMT) + mmt.is_multi_machine = True machine = MagicMock() machine.logs.return_value = "rank one" with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=MagicMock()), patch( "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=mmt, ), patch( - "lightning_sdk.cli.job.logs.resolve_mmt_machine", + "lightning_sdk.cli.job.logs.resolve_job_machine", return_value=machine, ) as resolve_rank: result = CliRunner().invoke(logs_job, ["distributed", "--rank", "1"]) @@ -100,7 +101,7 @@ def test_job_logs_selects_mmt_rank() -> None: assert "rank one" in result.output resolve_rank.assert_called_once_with(mmt, 1) machine.logs.assert_called_once_with( - follow=False, tail=None, rank=None, timestamps=False, since=None, until=None, query=None, severity=None + follow=False, tail=None, rank=0, timestamps=False, since=None, until=None, query=None, severity=None ) @@ -109,6 +110,7 @@ def test_job_logs_merges_mmt_without_rank() -> None: from lightning_sdk.cli.job.logs import logs_job mmt = MagicMock(spec=MMT) + mmt.is_multi_machine = True mmt.logs.return_value = "[distributed-0] zero\n[distributed-1] one" with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=MagicMock()), patch( "lightning_sdk.cli.job.logs.resolve_job_or_mmt", diff --git a/python/tests/cli/job/test_run.py b/python/tests/cli/job/test_run.py index ad45b2ab..243b3c48 100644 --- a/python/tests/cli/job/test_run.py +++ b/python/tests/cli/job/test_run.py @@ -70,23 +70,20 @@ def test_run_job_with_cloud(monkeypatch): @mock_command_logging -def test_run_job_with_multiple_machines_uses_mmt() -> None: +def test_run_job_with_multiple_machines_uses_job() -> None: submitted = MagicMock(name="distributed") submitted.name = "distributed" with patch("lightning_sdk.cli.job.run.resolve_teamspace", return_value=MagicMock()), patch( - "lightning_sdk.cli.job.run.Job.run" - ) as run_single, patch( - "lightning_sdk.cli.job.run.MMT.run", + "lightning_sdk.cli.job.run.Job.run", return_value=submitted, - ) as run_mmt: + ) as run_sdk_job: result = CliRunner().invoke( run_job, ["--name", "distributed", "--image", "ubuntu", "--num-machines", "3"], ) assert result.exit_code == 0, result.output - run_single.assert_not_called() - assert run_mmt.call_args.kwargs["num_machines"] == 3 + assert run_sdk_job.call_args.kwargs["num_machines"] == 3 assert "Submitted job distributed." in result.output diff --git a/python/tests/cli/job/test_ssh.py b/python/tests/cli/job/test_ssh.py index 0820bb8f..ba2b3c09 100644 --- a/python/tests/cli/job/test_ssh.py +++ b/python/tests/cli/job/test_ssh.py @@ -91,6 +91,7 @@ def test_ssh_runs_against_job_gateway_user() -> None: def test_ssh_selects_mmt_rank() -> None: mmt = MagicMock(spec=MMT) + mmt.is_multi_machine = True machine = MagicMock() machine.name = "distributed-1" machine.status = Status.Running @@ -100,7 +101,7 @@ def test_ssh_selects_mmt_rank() -> None: "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", return_value=mmt, ), patch( - "lightning_sdk.cli.job.ssh.resolve_mmt_machine", + "lightning_sdk.cli.job.ssh.resolve_job_machine", return_value=machine, ) as resolve_rank, patch( "lightning_sdk.cli.job.ssh.configure_ssh_internal", diff --git a/python/tests/cli/utils/test_resource_resolution.py b/python/tests/cli/utils/test_resource_resolution.py index 29337aac..280bfda2 100644 --- a/python/tests/cli/utils/test_resource_resolution.py +++ b/python/tests/cli/utils/test_resource_resolution.py @@ -8,6 +8,7 @@ join_teamspace_slug, resolve_cluster, resolve_job, + resolve_job_machine, resolve_job_or_mmt, resolve_mmt, resolve_mmt_machine, @@ -156,18 +157,15 @@ def test_resolve_job_converts_not_found_to_usage_error() -> None: resolve_job("train", MagicMock()) -def test_resolve_job_or_mmt_falls_back_to_mmt() -> None: +def test_resolve_job_or_mmt_uses_unified_job() -> None: teamspace = MagicMock() resolved = MagicMock() with patch( "lightning_sdk.cli.utils.resource_resolution.Job", - side_effect=ValueError("missing"), - ), patch( - "lightning_sdk.cli.utils.resource_resolution.MMT", return_value=resolved, - ) as mmt: + ) as job: assert resolve_job_or_mmt("distributed", teamspace) is resolved - mmt.assert_called_once_with(name="distributed", teamspace=teamspace) + job.assert_called_once_with(name="distributed", teamspace=teamspace) def test_resolve_mmt_machine_uses_rank_suffix() -> None: @@ -179,6 +177,7 @@ def test_resolve_mmt_machine_uses_rank_suffix() -> None: mmt.name = "distributed" assert resolve_mmt_machine(mmt, 1) is rank1 + assert resolve_job_machine(mmt, 1) is rank1 def test_resolve_mmt_requires_name() -> None: @@ -208,10 +207,10 @@ def test_resolve_mmt_converts_not_found_to_usage_error() -> None: def test_resolve_mmt_converts_real_lookup_not_found_to_usage_error() -> None: teamspace = MagicMock() teamspace.id = "teamspace-id" - with patch("lightning_sdk.mmt._resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.mmt.raise_access_error_if_not_allowed" - ), patch("lightning_sdk.mmt.CloudAccountApi"), patch( - "lightning_sdk.mmt.MMTApiV2.get_job_by_name", + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.raise_access_error_if_not_allowed" + ), patch("lightning_sdk.job.CloudAccountApi"), patch( + "lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404, reason="Not Found"), ), pytest.raises(click.UsageError, match="distributed"): resolve_mmt("distributed", teamspace) @@ -219,7 +218,7 @@ def test_resolve_mmt_converts_real_lookup_not_found_to_usage_error() -> None: def test_resolve_mmt_preserves_pre_lookup_not_found_error() -> None: failure = ApiException(status=404, reason="Teamspace Not Found") - with patch("lightning_sdk.mmt._resolve_teamspace", side_effect=failure), pytest.raises(ApiException) as raised: + with patch("lightning_sdk.job._resolve_teamspace", side_effect=failure), pytest.raises(ApiException) as raised: resolve_mmt("distributed", MagicMock()) assert raised.value is failure @@ -229,10 +228,10 @@ def test_resolve_mmt_preserves_real_lookup_permission_error() -> None: teamspace = MagicMock() teamspace.id = "teamspace-id" failure = ApiException(status=403, reason="Forbidden") - with patch("lightning_sdk.mmt._resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.mmt.raise_access_error_if_not_allowed" - ), patch("lightning_sdk.mmt.CloudAccountApi"), patch( - "lightning_sdk.mmt.MMTApiV2.get_job_by_name", + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.raise_access_error_if_not_allowed" + ), patch("lightning_sdk.job.CloudAccountApi"), patch( + "lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=failure, ), pytest.raises(ApiException) as raised: resolve_mmt("distributed", teamspace) diff --git a/python/tests/core/test_job.py b/python/tests/core/test_job.py index 25018fdb..d32046e7 100644 --- a/python/tests/core/test_job.py +++ b/python/tests/core/test_job.py @@ -11,6 +11,7 @@ V1Job, V1JobSpec, ) +from lightning_sdk.lightning_cloud.openapi.rest import ApiException from lightning_sdk.machine import Machine from lightning_sdk.status import Status from lightning_sdk.studio import Studio @@ -42,7 +43,9 @@ def init_job_error(): studio = Studio("st-abc", "ts-abc", org="org-abc") Job("xyz", studio.teamspace) - with pytest.raises(ValueError, match="Job xyz does not exist in Teamspace ts-abc"): + with mock.patch( + "lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404) + ), pytest.raises(ValueError, match="Job xyz does not exist in Teamspace ts-abc"): init_job_error() @@ -109,7 +112,9 @@ def test_delete_job( with pytest.raises(RuntimeError, match="Job j-abc does not exist in Teamspace ts-abc. Did you delete it?"): job.status # noqa: B018 - with pytest.raises(ValueError, match="Job j-abc does not exist in Teamspace ts-abc"): + with mock.patch( + "lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404) + ), pytest.raises(ValueError, match="Job j-abc does not exist in Teamspace ts-abc"): Job("j-abc", studio.teamspace) @@ -626,6 +631,7 @@ def test_submit_jobv2_studio_resolve( reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) @@ -724,6 +730,7 @@ def test_submit_job_v2_image_from_studio( reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) assert keeping_alive_mock.call_count == 0 @@ -769,6 +776,7 @@ def test_run_job_with_cloud_provider( reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) diff --git a/python/tests/core/test_permissions_integration.py b/python/tests/core/test_permissions_integration.py index 55c18e29..a8dc96b3 100644 --- a/python/tests/core/test_permissions_integration.py +++ b/python/tests/core/test_permissions_integration.py @@ -244,7 +244,7 @@ def test_job_init_raises_error_when_jobs_disabled(mock_teamspace_api, mock_resol # MMT class permission tests -@mock.patch("lightning_sdk.mmt._resolve_teamspace") +@mock.patch("lightning_sdk.job._resolve_teamspace") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @pytest.mark.project_permission_test() def test_mmt_init_raises_error_when_jobs_disabled(mock_teamspace_api, mock_resolve_teamspace): diff --git a/python/tests/core/test_pipeline.py b/python/tests/core/test_pipeline.py index 82203d1a..ecfa1ec5 100644 --- a/python/tests/core/test_pipeline.py +++ b/python/tests/core/test_pipeline.py @@ -171,6 +171,15 @@ def test_job_step_threads_placement_group_id(): assert proto.job.spec.placement_group_id == "pg-1" +@patch("lightning_sdk.pipeline.steps.CloudAccountApi", new=MagicMock()) +def test_job_step_supports_multiple_machines(): + job = JobStep(name="job-0", machine=Machine.CPU, num_machines=3) + proto = job.to_proto(MagicMock(), "", False) + + assert proto.type == V1PipelineStepType.MMT + assert proto.mmt.machines == 3 + + @pytest.mark.parametrize("interruption_retries", [0, 3]) @patch.object(teamspace, "TeamspaceApi", new=MagicMock()) @patch.object(pipeline_module, "_get_cluster", new=MagicMock()) diff --git a/python/tests/core/test_teamspace.py b/python/tests/core/test_teamspace.py index da06fa5f..28990a34 100644 --- a/python/tests/core/test_teamspace.py +++ b/python/tests/core/test_teamspace.py @@ -553,30 +553,39 @@ def test_download_model_version( ) +@mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_jobs") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) def test_list_jobs( list_jobs_mock, + list_mmts_mock, internal_get_org_api_mocker, internal_teamspace_api_mocker, internal_user_api_mocker, ): jobs = [V1Job(name="jobv2-1"), V1Job(name="jobv2-2"), V1Job(name="jobv2-3")] + mmts = [V1MultiMachineJob(name="mmtv2-1")] ts = Teamspace("ts-abc", org="org-abc") list_jobs_mock.return_value = jobs + list_mmts_mock.return_value = mmts # it's important that there are no additional calls to fetch individual jobs here. # they'd raise API Errors since we only mock the teamspace APIs listing # and not individual fetch requests listed_jobs = ts.jobs - assert len(listed_jobs) == 3 + assert len(listed_jobs) == 4 assert all(isinstance(j, Job) for j in listed_jobs) for lj, jj in zip(listed_jobs, jobs): assert lj.name == jj.name assert lj._job is jj + assert not lj.is_multi_machine + + assert listed_jobs[-1].name == "mmtv2-1" + assert listed_jobs[-1]._job is mmts[0] + assert listed_jobs[-1].is_multi_machine @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts") diff --git a/python/tests/core/test_unified_job.py b/python/tests/core/test_unified_job.py new file mode 100644 index 00000000..8bf87fc2 --- /dev/null +++ b/python/tests/core/test_unified_job.py @@ -0,0 +1,129 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from lightning_sdk.job import Job +from lightning_sdk.lightning_cloud.openapi import V1Job, V1JobSpec, V1MultiMachineJob +from lightning_sdk.lightning_cloud.openapi.rest import ApiException +from lightning_sdk.mmt import MMT + + +def _teamspace() -> SimpleNamespace: + return SimpleNamespace( + id="teamspace-id", + name="teamspace", + owner=SimpleNamespace(name="owner"), + default_cloud_account="default-cloud", + ) + + +def test_job_lookup_falls_back_to_multi_machine() -> None: + teamspace = _teamspace() + standalone_api = MagicMock() + standalone_api.get_job_by_name.side_effect = ApiException(status=404) + multi_api = MagicMock() + multi_api.get_job_by_name.return_value = V1MultiMachineJob( + id="mmt-id", name="distributed", machines=2, spec=V1JobSpec() + ) + multi_api.get_num_machines.return_value = 2 + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.JobApiV2", return_value=standalone_api + ), patch("lightning_sdk.job.MMTApiV2", return_value=multi_api): + job = Job("distributed", teamspace) + + assert job.is_multi_machine + assert job.num_machines == 2 + standalone_api.get_job_by_name.assert_called_once() + multi_api.get_job_by_name.assert_called_once() + + +def test_job_lookup_prefers_standalone_on_name_collision() -> None: + teamspace = _teamspace() + standalone_api = MagicMock() + standalone_api.get_job_by_name.return_value = V1Job(id="job-id", name="train", spec=V1JobSpec()) + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.JobApiV2", return_value=standalone_api + ), patch("lightning_sdk.job.MMTApiV2") as multi_api: + job = Job("train", teamspace) + + assert not job.is_multi_machine + assert job.num_machines == 1 + assert job.machines == (job,) + multi_api.assert_not_called() + + +def test_job_run_routes_multi_machine_submission() -> None: + teamspace = _teamspace() + multi_api = MagicMock() + multi_api.submit_job.return_value = V1MultiMachineJob( + id="mmt-id", name="distributed", machines=3, spec=V1JobSpec() + ) + multi_api.get_num_machines.return_value = 3 + cloud_api = MagicMock() + cloud_api.resolve_cloud_account.return_value = "cloud-id" + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job._resolve_default_cloud_account", return_value=None + ), patch("lightning_sdk.job.CloudAccountApi", return_value=cloud_api), patch( + "lightning_sdk.job.MMTApiV2", return_value=multi_api + ): + job = Job.run( + name="distributed", + machine="CPU", + cloud="aws", + image="ubuntu", + teamspace=teamspace, + num_machines=3, + ) + + assert type(job) is Job + assert job.is_multi_machine + assert job.num_machines == 3 + assert multi_api.submit_job.call_args.kwargs["num_machines"] == 3 + + +def test_mmt_run_returns_compatibility_subclass() -> None: + teamspace = _teamspace() + multi_api = MagicMock() + multi_api.submit_job.return_value = V1MultiMachineJob( + id="mmt-id", name="distributed", machines=2, spec=V1JobSpec() + ) + cloud_api = MagicMock() + cloud_api.resolve_cloud_account.return_value = "cloud-id" + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job._resolve_default_cloud_account", return_value=None + ), patch("lightning_sdk.job.CloudAccountApi", return_value=cloud_api), patch( + "lightning_sdk.job.MMTApiV2", return_value=multi_api + ): + job = MMT.run( + name="distributed", + num_machines=2, + machine="CPU", + cloud="aws", + image="ubuntu", + teamspace=teamspace, + ) + + assert isinstance(job, MMT) + assert isinstance(job, Job) + assert job.is_multi_machine + + +def test_mmt_lookup_skips_standalone_api() -> None: + teamspace = _teamspace() + multi_api = MagicMock() + multi_api.get_job_by_name.return_value = V1MultiMachineJob( + id="mmt-id", name="distributed", machines=2, spec=V1JobSpec() + ) + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.JobApiV2" + ) as standalone_api, patch("lightning_sdk.job.MMTApiV2", return_value=multi_api): + job = MMT("distributed", teamspace) + + assert isinstance(job, MMT) + assert job.is_multi_machine + standalone_api.return_value.get_job_by_name.assert_not_called() + multi_api.get_job_by_name.assert_called_once() From 55e8290145ecb658bda181e9f66b949c7edb2808 Mon Sep 17 00:00:00 2001 From: Teja Pulagam Date: Fri, 31 Jul 2026 19:10:12 +0000 Subject: [PATCH 3/6] fixes --- .../cli/utils/resource_resolution.py | 12 +++++-- python/lightning_sdk/job.py | 7 ++-- python/lightning_sdk/mmt.py | 34 +++++++++++++------ .../cli/utils/test_resource_resolution.py | 17 +++++++++- python/tests/core/test_job.py | 12 +++---- python/tests/core/test_unified_job.py | 30 ++++++++++++---- 6 files changed, 82 insertions(+), 30 deletions(-) diff --git a/python/lightning_sdk/cli/utils/resource_resolution.py b/python/lightning_sdk/cli/utils/resource_resolution.py index cfba72c2..bb8855dd 100644 --- a/python/lightning_sdk/cli/utils/resource_resolution.py +++ b/python/lightning_sdk/cli/utils/resource_resolution.py @@ -89,18 +89,26 @@ def resolve_job_machine(job: Job, rank: int) -> Job: if not machines: raise click.ClickException(f"Job '{job.name}' has no machines.") + for machine in machines: + if machine.rank == rank: + return machine + + # Fallback for older naming when rank is missing on the machine object. expected = f"{job.name}-{rank}" for machine in machines: if machine.name == expected: return machine + available_ranks: set[int] = set() prefix = f"{job.name}-" - available_ranks = [] for machine in machines: + if machine.rank is not None: + available_ranks.add(machine.rank) + continue if machine.name.startswith(prefix): suffix = machine.name[len(prefix) :] if suffix.isdigit(): - available_ranks.append(int(suffix)) + available_ranks.add(int(suffix)) available = ", ".join(str(value) for value in sorted(available_ranks)) raise click.ClickException(f"Rank {rank} not found on job '{job.name}'. Available ranks: {available or 'none'}.") diff --git a/python/lightning_sdk/job.py b/python/lightning_sdk/job.py index 109ba578..71406913 100644 --- a/python/lightning_sdk/job.py +++ b/python/lightning_sdk/job.py @@ -867,8 +867,7 @@ def _stream_logs( def link(self) -> str: if self.is_multi_machine: return ( - f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/" - f"jobs/{self.name}?app_id=mmt" + f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/" f"jobs/{self.name}?app_id=mmt" ) mmt_name = self._job_api.get_mmt_name(self._guaranteed_job) @@ -920,9 +919,7 @@ def _update_internal_job(self) -> None: from lightning_sdk.lightning_cloud.openapi.rest import ApiException try: - self._job = self._standalone_job_api.get_job_by_name( - name=self._name, teamspace_id=self._teamspace.id - ) + self._job = self._standalone_job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) except ApiException as ex: if ex.status != 404: raise diff --git a/python/lightning_sdk/mmt.py b/python/lightning_sdk/mmt.py index 3155d901..31df6b67 100644 --- a/python/lightning_sdk/mmt.py +++ b/python/lightning_sdk/mmt.py @@ -17,33 +17,43 @@ class MMTMachine(Protocol): """A single machine in a multi-machine job.""" @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def machine(self) -> Union["Machine", str]: ... + def machine(self) -> Union["Machine", str]: + ... @property - def artifact_path(self) -> Optional[str]: ... + def artifact_path(self) -> Optional[str]: + ... @property - def status(self) -> Status: ... + def status(self) -> Status: + ... @property - def resource_id(self) -> Optional[str]: ... + def resource_id(self) -> Optional[str]: + ... @property - def private_ip_address(self) -> Optional[str]: ... + def private_ip_address(self) -> Optional[str]: + ... @property - def placement_group_id(self) -> Optional[str]: ... + def placement_group_id(self) -> Optional[str]: + ... @property - def rank(self) -> Optional[int]: ... + def rank(self) -> Optional[int]: + ... @property - def logs(self) -> str: ... + def logs(self) -> str: + ... - def dict(self) -> JobDict: ... + def dict(self) -> JobDict: + ... class MMT(Job): @@ -71,6 +81,10 @@ def __init__( _resource_kind="multi", ) except ValueError as ex: + # Job.__init__ raises "Job {name} does not exist…" on 404; keep the MMT-specific + # wording for that case only. Propagate teamspace/validation errors unchanged. + if "does not exist in Teamspace" not in str(ex): + raise resolved_teamspace = getattr(self, "_teamspace", None) teamspace_name = getattr(resolved_teamspace, "name", teamspace) raise ValueError(f"Multi-machine job {name} does not exist in Teamspace {teamspace_name}") from ex diff --git a/python/tests/cli/utils/test_resource_resolution.py b/python/tests/cli/utils/test_resource_resolution.py index 34a3b116..06d20887 100644 --- a/python/tests/cli/utils/test_resource_resolution.py +++ b/python/tests/cli/utils/test_resource_resolution.py @@ -198,11 +198,26 @@ def test_resolve_job_or_mmt_uses_unified_job() -> None: job.assert_called_once_with(name="distributed", teamspace=teamspace) -def test_resolve_mmt_machine_uses_rank_suffix() -> None: +def test_resolve_mmt_machine_prefers_rank_attribute() -> None: + rank0 = MagicMock(name="odd-name-0") + rank0.name = "odd-name-0" + rank0.rank = 0 + rank1 = MagicMock(name="odd-name-1") + rank1.name = "odd-name-1" + rank1.rank = 1 + mmt = MagicMock(name="distributed", machines=(rank0, rank1)) + mmt.name = "distributed" + + assert resolve_job_machine(mmt, 1) is rank1 + + +def test_resolve_mmt_machine_falls_back_to_name_suffix() -> None: rank0 = MagicMock(name="rank-0") rank0.name = "distributed-0" + rank0.rank = None rank1 = MagicMock(name="rank-1") rank1.name = "distributed-1" + rank1.rank = None mmt = MagicMock(name="distributed", machines=(rank0, rank1)) mmt.name = "distributed" diff --git a/python/tests/core/test_job.py b/python/tests/core/test_job.py index d32046e7..268082eb 100644 --- a/python/tests/core/test_job.py +++ b/python/tests/core/test_job.py @@ -43,9 +43,9 @@ def init_job_error(): studio = Studio("st-abc", "ts-abc", org="org-abc") Job("xyz", studio.teamspace) - with mock.patch( - "lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404) - ), pytest.raises(ValueError, match="Job xyz does not exist in Teamspace ts-abc"): + with mock.patch("lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404)), pytest.raises( + ValueError, match="Job xyz does not exist in Teamspace ts-abc" + ): init_job_error() @@ -112,9 +112,9 @@ def test_delete_job( with pytest.raises(RuntimeError, match="Job j-abc does not exist in Teamspace ts-abc. Did you delete it?"): job.status # noqa: B018 - with mock.patch( - "lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404) - ), pytest.raises(ValueError, match="Job j-abc does not exist in Teamspace ts-abc"): + with mock.patch("lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404)), pytest.raises( + ValueError, match="Job j-abc does not exist in Teamspace ts-abc" + ): Job("j-abc", studio.teamspace) diff --git a/python/tests/core/test_unified_job.py b/python/tests/core/test_unified_job.py index 8bf87fc2..ccda2162 100644 --- a/python/tests/core/test_unified_job.py +++ b/python/tests/core/test_unified_job.py @@ -1,6 +1,8 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch +import pytest + from lightning_sdk.job import Job from lightning_sdk.lightning_cloud.openapi import V1Job, V1JobSpec, V1MultiMachineJob from lightning_sdk.lightning_cloud.openapi.rest import ApiException @@ -56,9 +58,7 @@ def test_job_lookup_prefers_standalone_on_name_collision() -> None: def test_job_run_routes_multi_machine_submission() -> None: teamspace = _teamspace() multi_api = MagicMock() - multi_api.submit_job.return_value = V1MultiMachineJob( - id="mmt-id", name="distributed", machines=3, spec=V1JobSpec() - ) + multi_api.submit_job.return_value = V1MultiMachineJob(id="mmt-id", name="distributed", machines=3, spec=V1JobSpec()) multi_api.get_num_machines.return_value = 3 cloud_api = MagicMock() cloud_api.resolve_cloud_account.return_value = "cloud-id" @@ -86,9 +86,7 @@ def test_job_run_routes_multi_machine_submission() -> None: def test_mmt_run_returns_compatibility_subclass() -> None: teamspace = _teamspace() multi_api = MagicMock() - multi_api.submit_job.return_value = V1MultiMachineJob( - id="mmt-id", name="distributed", machines=2, spec=V1JobSpec() - ) + multi_api.submit_job.return_value = V1MultiMachineJob(id="mmt-id", name="distributed", machines=2, spec=V1JobSpec()) cloud_api = MagicMock() cloud_api.resolve_cloud_account.return_value = "cloud-id" @@ -127,3 +125,23 @@ def test_mmt_lookup_skips_standalone_api() -> None: assert job.is_multi_machine standalone_api.return_value.get_job_by_name.assert_not_called() multi_api.get_job_by_name.assert_called_once() + + +def test_mmt_lookup_rewrites_only_not_found_errors() -> None: + teamspace = _teamspace() + multi_api = MagicMock() + multi_api.get_job_by_name.side_effect = ApiException(status=404) + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.JobApiV2" + ), patch("lightning_sdk.job.MMTApiV2", return_value=multi_api), pytest.raises( + ValueError, match="Multi-machine job distributed does not exist in Teamspace teamspace" + ): + MMT("distributed", teamspace) + + +def test_mmt_lookup_propagates_teamspace_resolution_errors() -> None: + with patch("lightning_sdk.job._resolve_teamspace", return_value=None), pytest.raises( + ValueError, match="Cannot resolve the teamspace" + ): + MMT("distributed", "owner/teamspace") From e4ac9b7f85337282fcd63dc175a1908b59cb515d Mon Sep 17 00:00:00 2001 From: Teja Pulagam Date: Fri, 31 Jul 2026 19:24:35 +0000 Subject: [PATCH 4/6] fix --- python/lightning_sdk/cli/mmt/ssh.py | 30 ++--------------------- python/lightning_sdk/job.py | 4 +--- python/tests/cli/mmt/test_ssh.py | 37 ++++++++++++++++++++++++++++- python/tests/core/test_job.py | 34 +++++++++++++------------- 4 files changed, 55 insertions(+), 50 deletions(-) diff --git a/python/lightning_sdk/cli/mmt/ssh.py b/python/lightning_sdk/cli/mmt/ssh.py index 747d2491..ad5cd547 100644 --- a/python/lightning_sdk/cli/mmt/ssh.py +++ b/python/lightning_sdk/cli/mmt/ssh.py @@ -6,10 +6,8 @@ import rich_click as click from lightning_sdk.cli.utils.logging import LightningCommand -from lightning_sdk.cli.utils.resource_resolution import resolve_mmt, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job_machine, resolve_mmt, resolve_teamspace from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal -from lightning_sdk.job import Job -from lightning_sdk.mmt import MMT from lightning_sdk.status import Status _SSH_HOST = "ssh.lightning.ai" @@ -22,30 +20,6 @@ def _ssh_user_for_job_id(job_id: str) -> str: return f"j_{suffix}" -def _machine_for_rank(mmt: MMT, rank: int) -> Job: - machines = mmt.machines - if not machines: - raise click.ClickException(f"Multi-machine job '{mmt.name}' has no machines to SSH into.") - - expected = f"{mmt.name}-{rank}" - for machine in machines: - if machine.name == expected: - return machine - - prefix = f"{mmt.name}-" - available_ranks = [] - for machine in machines: - if not machine.name.startswith(prefix): - continue - suffix = machine.name[len(prefix) :] - if suffix.isdigit(): - available_ranks.append(int(suffix)) - available = ", ".join(str(r) for r in sorted(available_ranks)) - raise click.ClickException( - f"Rank {rank} not found on multi-machine job '{mmt.name}'. Available ranks: {available or 'none'}." - ) - - @click.command("ssh", cls=LightningCommand) @click.argument("name") @click.option( @@ -84,7 +58,7 @@ def ssh_impl( ) -> None: resolved_teamspace = resolve_teamspace(teamspace) mmt = resolve_mmt(name, resolved_teamspace) - job = _machine_for_rank(mmt, rank) + job = resolve_job_machine(mmt, rank) if job.status != Status.Running: raise click.ClickException( diff --git a/python/lightning_sdk/job.py b/python/lightning_sdk/job.py index 71406913..02e148a9 100644 --- a/python/lightning_sdk/job.py +++ b/python/lightning_sdk/job.py @@ -866,9 +866,7 @@ def _stream_logs( @property def link(self) -> str: if self.is_multi_machine: - return ( - f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/" f"jobs/{self.name}?app_id=mmt" - ) + return f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=mmt" mmt_name = self._job_api.get_mmt_name(self._guaranteed_job) diff --git a/python/tests/cli/mmt/test_ssh.py b/python/tests/cli/mmt/test_ssh.py index cc3697c6..a15105ef 100644 --- a/python/tests/cli/mmt/test_ssh.py +++ b/python/tests/cli/mmt/test_ssh.py @@ -60,6 +60,7 @@ def test_ssh_resolves_before_downloading_keys() -> None: def test_ssh_rejects_non_running_machine() -> None: rank0 = MagicMock() rank0.name = "train-0" + rank0.rank = 0 rank0.status = Status.Completed rank0.id = "job_01abc" @@ -81,11 +82,13 @@ def test_ssh_defaults_to_rank_zero() -> None: teamspace = MagicMock() rank0 = MagicMock() rank0.name = "train-0" + rank0.rank = 0 rank0.status = Status.Running rank0.id = "job_rank0" rank1 = MagicMock() rank1.name = "train-1" + rank1.rank = 1 rank1.status = Status.Running rank1.id = "job_rank1" @@ -109,11 +112,13 @@ def test_ssh_selects_requested_rank() -> None: teamspace = MagicMock() rank0 = MagicMock() rank0.name = "olmo3-7b-think-sft-full-0" + rank0.rank = 0 rank0.status = Status.Running rank0.id = "job_r0" rank1 = MagicMock() rank1.name = "olmo3-7b-think-sft-full-1" + rank1.rank = 1 rank1.status = Status.Running rank1.id = "job_r1" @@ -135,6 +140,7 @@ def test_ssh_selects_requested_rank() -> None: def test_ssh_rejects_unknown_rank() -> None: rank0 = MagicMock() rank0.name = "train-0" + rank0.rank = 0 mmt = MagicMock() mmt.name = "train" @@ -142,13 +148,42 @@ def test_ssh_rejects_unknown_rank() -> None: with patch("lightning_sdk.cli.mmt.ssh.resolve_teamspace", return_value=MagicMock()), patch( "lightning_sdk.cli.mmt.ssh.resolve_mmt", return_value=mmt - ), pytest.raises(click.ClickException, match="Rank 3 not found"): + ), pytest.raises(click.ClickException, match="Rank 3 not found.*Available ranks: 0"): ssh_impl(name="train", teamspace=None, rank=3) +def test_ssh_selects_rank_when_names_lack_rank_suffix() -> None: + """Machines are matched on their rank, not on a ``{job}-{rank}`` name convention.""" + rank0 = MagicMock() + rank0.name = "worker-alpha" + rank0.rank = 0 + rank0.status = Status.Running + rank0.id = "job_r0" + + rank1 = MagicMock() + rank1.name = "worker-beta" + rank1.rank = 1 + rank1.status = Status.Running + rank1.id = "job_r1" + + mmt = MagicMock() + mmt.name = "train" + mmt.machines = (rank0, rank1) + + with patch("lightning_sdk.cli.mmt.ssh.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.mmt.ssh.resolve_mmt", return_value=mmt + ), patch("lightning_sdk.cli.mmt.ssh.configure_ssh_internal", return_value="/tmp/lightning_rsa"), patch( + "lightning_sdk.cli.mmt.ssh.subprocess.run" + ) as run: + ssh_impl(name="train", teamspace=None, rank=1) + + run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_r1@ssh.lightning.ai"]) + + def test_ssh_retries_with_fresh_keys_on_failure() -> None: rank0 = MagicMock() rank0.name = "train-0" + rank0.rank = 0 rank0.status = Status.Running rank0.id = "job_01abc" diff --git a/python/tests/core/test_job.py b/python/tests/core/test_job.py index 268082eb..4e074d30 100644 --- a/python/tests/core/test_job.py +++ b/python/tests/core/test_job.py @@ -610,9 +610,10 @@ def test_submit_jobv2_studio_resolve( from lightning_sdk.job import Job submit_mock = mock.MagicMock() - Job._submit = submit_mock - - Job.run("test-job", machine=Machine.CPU, command="echo hello", studio="st-abc", teamspace="ts-abc", org="org-abc") + with mock.patch.object(Job, "_submit", submit_mock): + Job.run( + "test-job", machine=Machine.CPU, command="echo hello", studio="st-abc", teamspace="ts-abc", org="org-abc" + ) submit_mock.assert_called_once_with( command="echo hello", @@ -691,11 +692,10 @@ def test_submit_job_v2_image_from_studio( from lightning_sdk.job import Job submit_mock = mock.MagicMock() - Job._submit = submit_mock keeping_alive_mock = mock.MagicMock() StudioApi.start_keeping_alive = keeping_alive_mock - with mock.patch.dict( + with mock.patch.object(Job, "_submit", submit_mock), mock.patch.dict( os.environ, { "LIGHTNING_CLOUD_SPACE_ID": "st-abc", @@ -747,17 +747,16 @@ def test_run_job_with_cloud_provider( from lightning_sdk.job import Job submit_mock = mock.MagicMock() - Job._submit = submit_mock - - Job.run( - "test-job", - machine=Machine.CPU, - command="echo hello", - image="nginx", - teamspace="ts-abc", - org="org-abc", - cloud="nebius", - ) + with mock.patch.object(Job, "_submit", submit_mock): + Job.run( + "test-job", + machine=Machine.CPU, + command="echo hello", + image="nginx", + teamspace="ts-abc", + org="org-abc", + cloud="nebius", + ) submit_mock.assert_called_once_with( command="echo hello", @@ -952,11 +951,10 @@ def test_submit_job_from_running_studio( from lightning_sdk.job import Job submit_mock = mock.MagicMock() - Job._submit = submit_mock keeping_alive_mock = mock.MagicMock() StudioApi.start_keeping_alive = keeping_alive_mock - with mock.patch.dict( + with mock.patch.object(Job, "_submit", submit_mock), mock.patch.dict( os.environ, { "LIGHTNING_CLOUD_SPACE_ID": "st-abc", From 5c95d873a1c6e85cc1d2fdbd0ed0ec0414e22f1b Mon Sep 17 00:00:00 2001 From: Teja Pulagam Date: Fri, 31 Jul 2026 23:16:35 +0000 Subject: [PATCH 5/6] remove excess --- python/lightning_sdk/cli/job/inspect.py | 4 ++-- python/lightning_sdk/cli/job/logs.py | 4 ++-- python/lightning_sdk/cli/job/ssh.py | 4 ++-- python/lightning_sdk/cli/job/stop.py | 4 ++-- .../cli/utils/resource_resolution.py | 10 ---------- python/tests/cli/job/test_inspect.py | 4 ++-- python/tests/cli/job/test_logs.py | 14 +++++++------- python/tests/cli/job/test_ssh.py | 10 +++++----- python/tests/cli/job/test_stop.py | 2 +- .../cli/utils/test_resource_resolution.py | 18 ++---------------- 10 files changed, 25 insertions(+), 49 deletions(-) diff --git a/python/lightning_sdk/cli/job/inspect.py b/python/lightning_sdk/cli/job/inspect.py index b329ad47..368c24a1 100644 --- a/python/lightning_sdk/cli/job/inspect.py +++ b/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_machine, resolve_job_or_mmt, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace @click.command("inspect", cls=LightningCommand) @@ -30,7 +30,7 @@ def inspect_job( ) -> None: """Inspect a job for further details as JSON.""" resolved_teamspace = resolve_teamspace(teamspace) - job = resolve_job_or_mmt(name, resolved_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: diff --git a/python/lightning_sdk/cli/job/logs.py b/python/lightning_sdk/cli/job/logs.py index b7f4481e..530fd202 100644 --- a/python/lightning_sdk/cli/job/logs.py +++ b/python/lightning_sdk/cli/job/logs.py @@ -8,7 +8,7 @@ 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_machine, resolve_job_or_mmt, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace @click.command("logs", cls=LightningCommand) @@ -55,7 +55,7 @@ def logs_job( --rank). """ resolved_teamspace = resolve_teamspace(teamspace) - resource = resolve_job_or_mmt(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 diff --git a/python/lightning_sdk/cli/job/ssh.py b/python/lightning_sdk/cli/job/ssh.py index 768576fc..ae39ee22 100644 --- a/python/lightning_sdk/cli/job/ssh.py +++ b/python/lightning_sdk/cli/job/ssh.py @@ -6,7 +6,7 @@ import rich_click as click from lightning_sdk.cli.utils.logging import LightningCommand -from lightning_sdk.cli.utils.resource_resolution import resolve_job_machine, resolve_job_or_mmt, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal from lightning_sdk.status import Status @@ -54,7 +54,7 @@ def ssh_impl( raise click.UsageError("Missing job name. Pass NAME.") resolved_teamspace = resolve_teamspace(teamspace) - job = resolve_job_or_mmt(name, resolved_teamspace) + job = resolve_job(name, resolved_teamspace) if job.is_multi_machine is True: job = resolve_job_machine(job, rank if rank is not None else 0) elif rank is not None: diff --git a/python/lightning_sdk/cli/job/stop.py b/python/lightning_sdk/cli/job/stop.py index 1632e9a5..3c4adcec 100644 --- a/python/lightning_sdk/cli/job/stop.py +++ b/python/lightning_sdk/cli/job/stop.py @@ -7,7 +7,7 @@ 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_job_or_mmt, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_teamspace @click.command("stop", cls=LightningCommand) @@ -25,7 +25,7 @@ def stop_job(name: str, teamspace: Optional[str] = None, as_json: bool = False) -> None: """Stop a job.""" resolved_teamspace = resolve_teamspace(teamspace) - job = resolve_job_or_mmt(name, resolved_teamspace) + job = resolve_job(name, resolved_teamspace) job.stop() if as_json: echo_json({"name": job.name, "status": "stopped"}) diff --git a/python/lightning_sdk/cli/utils/resource_resolution.py b/python/lightning_sdk/cli/utils/resource_resolution.py index bb8855dd..7d05c384 100644 --- a/python/lightning_sdk/cli/utils/resource_resolution.py +++ b/python/lightning_sdk/cli/utils/resource_resolution.py @@ -78,11 +78,6 @@ def resolve_job(name: Optional[str], teamspace: Teamspace) -> Job: raise click.UsageError(f"Could not resolve job '{name}' in teamspace '{teamspace.name}'.") from ex -def resolve_job_or_mmt(name: Optional[str], teamspace: Teamspace) -> Job: - """Resolve a single- or multi-machine job by name.""" - return resolve_job(name, teamspace) - - def resolve_job_machine(job: Job, rank: int) -> Job: """Resolve one machine in a multi-machine job by rank.""" machines = job.machines @@ -113,11 +108,6 @@ def resolve_job_machine(job: Job, rank: int) -> Job: raise click.ClickException(f"Rank {rank} not found on job '{job.name}'. Available ranks: {available or 'none'}.") -def resolve_mmt_machine(mmt: MMT, rank: int) -> Job: - """Compatibility alias for resolving a machine by rank.""" - return resolve_job_machine(mmt, rank) - - def resolve_mmt(name: Optional[str], teamspace: Teamspace) -> MMT: if not name: raise click.UsageError("Missing multi-machine job name. Pass JOB.") diff --git a/python/tests/cli/job/test_inspect.py b/python/tests/cli/job/test_inspect.py index af294994..559f4ea8 100644 --- a/python/tests/cli/job/test_inspect.py +++ b/python/tests/cli/job/test_inspect.py @@ -55,7 +55,7 @@ def test_job_inspect_uses_positional_name() -> None: job.json.return_value = '{"name":"my-job"}' with patch("lightning_sdk.cli.job.inspect.resolve_teamspace", return_value=teamspace) as resolve_teamspace, patch( - "lightning_sdk.cli.job.inspect.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.inspect.resolve_job", return_value=job ) as resolve_job: result = CliRunner().invoke(inspect_job, ["my-job", "--teamspace", "org/teamspace"]) @@ -76,7 +76,7 @@ def test_job_inspect_selects_mmt_rank() -> None: machine = MagicMock() machine.json.return_value = '{"name":"distributed-1"}' with patch("lightning_sdk.cli.job.inspect.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.inspect.resolve_job_or_mmt", + "lightning_sdk.cli.job.inspect.resolve_job", return_value=mmt, ), patch( "lightning_sdk.cli.job.inspect.resolve_job_machine", diff --git a/python/tests/cli/job/test_logs.py b/python/tests/cli/job/test_logs.py index 7839bd2d..34578441 100644 --- a/python/tests/cli/job/test_logs.py +++ b/python/tests/cli/job/test_logs.py @@ -44,7 +44,7 @@ def test_job_logs_prints_snapshot() -> None: job = MagicMock() job.logs.return_value = "hello from the job\n42" with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=teamspace) as resolve_teamspace, patch( - "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.logs.resolve_job", return_value=job ) as resolve_job: result = CliRunner().invoke(logs_job, ["my-job", "--teamspace", "org/teamspace"]) @@ -66,7 +66,7 @@ def test_job_logs_follows_with_options() -> None: job = MagicMock() job.logs.return_value = iter(["line 1", "line 2"]) with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.logs.resolve_job", return_value=job ): result = CliRunner().invoke( logs_job, @@ -89,7 +89,7 @@ def test_job_logs_selects_mmt_rank() -> None: machine = MagicMock() machine.logs.return_value = "rank one" with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=MagicMock()), patch( - "lightning_sdk.cli.job.logs.resolve_job_or_mmt", + "lightning_sdk.cli.job.logs.resolve_job", return_value=mmt, ), patch( "lightning_sdk.cli.job.logs.resolve_job_machine", @@ -113,7 +113,7 @@ def test_job_logs_merges_mmt_without_rank() -> None: mmt.is_multi_machine = True mmt.logs.return_value = "[distributed-0] zero\n[distributed-1] one" with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=MagicMock()), patch( - "lightning_sdk.cli.job.logs.resolve_job_or_mmt", + "lightning_sdk.cli.job.logs.resolve_job", return_value=mmt, ): result = CliRunner().invoke(logs_job, ["distributed"]) @@ -134,7 +134,7 @@ def test_job_logs_passes_filters() -> None: job = MagicMock() job.logs.return_value = "boom" with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.logs.resolve_job", return_value=job ): result = CliRunner().invoke(logs_job, ["my-job", "--query", "boom", "--severity", "error"]) @@ -148,7 +148,7 @@ def test_job_logs_passes_filters() -> None: def test_job_logs_rejects_unknown_severity() -> None: from lightning_sdk.cli.job.logs import logs_job - with patch("lightning_sdk.cli.job.logs.resolve_job_or_mmt") as resolve_job: + with patch("lightning_sdk.cli.job.logs.resolve_job") as resolve_job: result = CliRunner().invoke(logs_job, ["my-job", "--severity", "critical"]) assert result.exit_code != 0 @@ -163,7 +163,7 @@ def test_job_logs_reports_sdk_errors_cleanly() -> None: job = MagicMock() job.logs.side_effect = RuntimeError("Logs are not available while the job is Pending.") with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.logs.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.logs.resolve_job", return_value=job ): result = CliRunner().invoke(logs_job, ["my-job"]) diff --git a/python/tests/cli/job/test_ssh.py b/python/tests/cli/job/test_ssh.py index ba2b3c09..cc09d637 100644 --- a/python/tests/cli/job/test_ssh.py +++ b/python/tests/cli/job/test_ssh.py @@ -43,7 +43,7 @@ def test_ssh_resolves_before_downloading_keys() -> None: "lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock(name="coding-model-training"), ), patch( - "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", + "lightning_sdk.cli.job.ssh.resolve_job", side_effect=click.UsageError("Could not resolve job 'missing'."), ), patch( "lightning_sdk.cli.job.ssh.configure_ssh_internal", @@ -61,7 +61,7 @@ def test_ssh_rejects_non_running_job() -> None: job.id = "job_01abc" with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch( - "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.ssh.resolve_job", return_value=job ), patch("lightning_sdk.cli.job.ssh.configure_ssh_internal") as configure, pytest.raises( click.ClickException, match="not Running" ): @@ -78,7 +78,7 @@ def test_ssh_runs_against_job_gateway_user() -> None: job.id = "job_01jj4hvvjj4zx1t1esm5az3zt7" with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.ssh.resolve_job", return_value=job ) as resolve_job, patch( "lightning_sdk.cli.job.ssh.configure_ssh_internal", return_value="/tmp/lightning_rsa" ), patch("lightning_sdk.cli.job.ssh.subprocess.run") as run: @@ -98,7 +98,7 @@ def test_ssh_selects_mmt_rank() -> None: machine.id = "job_rank1" with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch( - "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", + "lightning_sdk.cli.job.ssh.resolve_job", return_value=mmt, ), patch( "lightning_sdk.cli.job.ssh.resolve_job_machine", @@ -123,7 +123,7 @@ def test_ssh_retries_with_fresh_keys_on_failure() -> None: run = MagicMock(side_effect=[OSError("ssh missing"), None]) with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch( - "lightning_sdk.cli.job.ssh.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.ssh.resolve_job", return_value=job ), patch("lightning_sdk.cli.job.ssh.configure_ssh_internal", configure), patch( "lightning_sdk.cli.job.ssh.subprocess.run", run ): diff --git a/python/tests/cli/job/test_stop.py b/python/tests/cli/job/test_stop.py index 9904bdf7..3bacf48a 100644 --- a/python/tests/cli/job/test_stop.py +++ b/python/tests/cli/job/test_stop.py @@ -47,7 +47,7 @@ def test_job_stop_resolves_exact_name() -> None: job = MagicMock() job.name = "train" with patch("lightning_sdk.cli.job.stop.resolve_teamspace", return_value=teamspace) as resolve_teamspace, patch( - "lightning_sdk.cli.job.stop.resolve_job_or_mmt", return_value=job + "lightning_sdk.cli.job.stop.resolve_job", return_value=job ) as resolve_job: result = CliRunner().invoke(stop_job, ["train", "--teamspace", "org/teamspace"]) diff --git a/python/tests/cli/utils/test_resource_resolution.py b/python/tests/cli/utils/test_resource_resolution.py index 06d20887..2c024000 100644 --- a/python/tests/cli/utils/test_resource_resolution.py +++ b/python/tests/cli/utils/test_resource_resolution.py @@ -10,9 +10,7 @@ resolve_deployment, resolve_job, resolve_job_machine, - resolve_job_or_mmt, resolve_mmt, - resolve_mmt_machine, resolve_studio, resolve_teamspace, ) @@ -187,18 +185,7 @@ def test_resolve_job_converts_not_found_to_usage_error() -> None: resolve_job("train", MagicMock()) -def test_resolve_job_or_mmt_uses_unified_job() -> None: - teamspace = MagicMock() - resolved = MagicMock() - with patch( - "lightning_sdk.cli.utils.resource_resolution.Job", - return_value=resolved, - ) as job: - assert resolve_job_or_mmt("distributed", teamspace) is resolved - job.assert_called_once_with(name="distributed", teamspace=teamspace) - - -def test_resolve_mmt_machine_prefers_rank_attribute() -> None: +def test_resolve_job_machine_prefers_rank_attribute() -> None: rank0 = MagicMock(name="odd-name-0") rank0.name = "odd-name-0" rank0.rank = 0 @@ -211,7 +198,7 @@ def test_resolve_mmt_machine_prefers_rank_attribute() -> None: assert resolve_job_machine(mmt, 1) is rank1 -def test_resolve_mmt_machine_falls_back_to_name_suffix() -> None: +def test_resolve_job_machine_falls_back_to_name_suffix() -> None: rank0 = MagicMock(name="rank-0") rank0.name = "distributed-0" rank0.rank = None @@ -221,7 +208,6 @@ def test_resolve_mmt_machine_falls_back_to_name_suffix() -> None: mmt = MagicMock(name="distributed", machines=(rank0, rank1)) mmt.name = "distributed" - assert resolve_mmt_machine(mmt, 1) is rank1 assert resolve_job_machine(mmt, 1) is rank1 From 79f2fb70e0ce29a5c6b4a398a04ba2f6813eb9de Mon Sep 17 00:00:00 2001 From: Teja Pulagam Date: Mon, 3 Aug 2026 18:38:04 +0000 Subject: [PATCH 6/6] unify public interfaces + detect mmt with num machines --- python/lightning_sdk/api/job_api.py | 17 +++- python/lightning_sdk/api/mmt_api.py | 7 +- python/lightning_sdk/job.py | 127 +++++++++++++------------- python/lightning_sdk/mmt.py | 4 +- python/lightning_sdk/teamspace.py | 8 +- python/tests/core/test_job.py | 2 + python/tests/core/test_mmt.py | 4 +- python/tests/core/test_unified_job.py | 2 +- 8 files changed, 98 insertions(+), 73 deletions(-) diff --git a/python/lightning_sdk/api/job_api.py b/python/lightning_sdk/api/job_api.py index 12c0b7ea..b0c4da84 100644 --- a/python/lightning_sdk/api/job_api.py +++ b/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 [] diff --git a/python/lightning_sdk/api/mmt_api.py b/python/lightning_sdk/api/mmt_api.py index ad6c4410..1c0d934a 100644 --- a/python/lightning_sdk/api/mmt_api.py +++ b/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) diff --git a/python/lightning_sdk/job.py b/python/lightning_sdk/job.py index 02e148a9..f6a5a5b7 100644 --- a/python/lightning_sdk/job.py +++ b/python/lightning_sdk/job.py @@ -132,7 +132,7 @@ def __init__( user: Union[str, "User", None] = None, *, _fetch_job: bool = True, - _resource_kind: Optional[str] = None, + _num_machines: int = 1, ) -> None: """Fetch already existing jobs. @@ -150,6 +150,9 @@ def __init__( when ``_fetch_job=True``. PermissionError: If the user does not have access to jobs in the given teamspace. """ + if _num_machines < 1: + raise ValueError("A job needs to run on at least one machine") + teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user) if teamspace is None: raise ValueError( @@ -165,11 +168,8 @@ def __init__( self._prevent_refetch_latest = False self._cloud_account_api = CloudAccountApi() self._standalone_job_api = JobApiV2() - self._mmt_job_api: Optional[MMTApiV2] = None - self._resource_kind = _resource_kind or "standalone" - self._job_api: Union[JobApiV2, MMTApiV2] = self._standalone_job_api - if self._resource_kind == "multi": - self._set_resource_kind("multi") + self._mmt_job_api = MMTApiV2() + self._num_machines = _num_machines self._logs_api = LogsApi() if _fetch_job: @@ -182,16 +182,20 @@ def __init__( raise ValueError(f"Job {name} does not exist in Teamspace {teamspace.name}") from None raise - def _set_resource_kind(self, kind: str) -> None: - if kind not in ("standalone", "multi"): - raise ValueError(f"Unknown job resource kind: {kind}") - self._resource_kind = kind - if kind == "multi": - if self._mmt_job_api is None: - self._mmt_job_api = MMTApiV2() - self._job_api = self._mmt_job_api - else: - self._job_api = self._standalone_job_api + @property + def _job_api(self) -> Union[JobApiV2, MMTApiV2]: + return self._mmt_job_api if self._num_machines > 1 else self._standalone_job_api + + def _attach_job(self, job: Any) -> None: + """Bind a fetched job payload and sync ``num_machines`` from it.""" + from lightning_sdk.lightning_cloud.openapi import V1Job, V1MultiMachineJob + + self._job = job + # Only a known payload type changes the machine count; anything else keeps the current one. + if isinstance(job, V1MultiMachineJob): + self._num_machines = job.machines if job.machines and job.machines > 1 else max(self._num_machines, 2) + elif isinstance(job, V1Job): + self._num_machines = 1 @classmethod def run( @@ -331,8 +335,7 @@ def run( elif entrypoint == "" or entrypoint is None: entrypoint = None - job = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False) - job._set_resource_kind("multi" if num_machines > 1 else "standalone") + job = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False, _num_machines=num_machines) submit_cloud = cloud if cloud_account is None else None job._submit( @@ -380,9 +383,7 @@ def _submit( ) -> "Job": if num_machines < 1: raise ValueError("A job needs to run on at least one machine") - if self.is_multi_machine and num_machines <= 1: - raise ValueError("Multi-machine jobs need to run on at least two machines") - if self.is_multi_machine and scratch_disks: + if num_machines > 1 and scratch_disks: raise ValueError("scratch_disks are not supported for multi-machine jobs") if studio is not None: @@ -425,36 +426,35 @@ def _submit( if ".." in path.parts: raise ValueError("scratch_disk path cannot contain '..'") - submit_kwargs = { - "name": self.name, - "command": command, - "cloud_account": cloud_account, - "teamspace_id": self._teamspace.id, - "studio_id": studio_id, - "image": image, - "machine": machine, - "interruptible": interruptible, - "env": env, - "image_credentials": image_credentials, - "cloud_account_auth": cloud_account_auth, - "entrypoint": entrypoint, - "path_mappings": path_mappings, - "max_runtime": max_runtime, - "reuse_snapshot": reuse_snapshot, - "placement_group_id": placement_group_id, - } - if self.is_multi_machine: - submitted = self._job_api.submit_job(num_machines=num_machines, **submit_kwargs) - else: - submitted = self._job_api.submit_job(scratch_disks=scratch_disks, **submit_kwargs) - if not self.is_multi_machine and submitted.name != self._name: + self._num_machines = num_machines + submitted = self._job_api.submit_job( + name=self.name, + command=command, + cloud_account=cloud_account, + teamspace_id=self._teamspace.id, + studio_id=studio_id, + image=image, + machine=machine, + interruptible=interruptible, + env=env, + image_credentials=image_credentials, + cloud_account_auth=cloud_account_auth, + entrypoint=entrypoint, + path_mappings=path_mappings, + max_runtime=max_runtime, + reuse_snapshot=reuse_snapshot, + placement_group_id=placement_group_id, + num_machines=num_machines, + scratch_disks=scratch_disks, + ) + if num_machines <= 1 and submitted.name != self._name: warnings.warn( f"Job name '{self._name}' was already taken in this teamspace; " f"the job was created as '{submitted.name}' instead.", stacklevel=2, ) - self._job = submitted + self._attach_job(submitted) self._name = submitted.name return self @@ -465,14 +465,12 @@ def stop(self) -> None: self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id) def delete(self) -> None: - if self.is_multi_machine: - self._job_api.delete_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id) - else: - self._job_api.delete_job( - job_id=self._guaranteed_job.id, - teamspace_id=self._teamspace.id, - cloudspace_id=self._guaranteed_job.spec.cloudspace_id, - ) + cloudspace_id = None if self.is_multi_machine else self._guaranteed_job.spec.cloudspace_id + self._job_api.delete_job( + job_id=self._guaranteed_job.id, + teamspace_id=self._teamspace.id, + cloudspace_id=cloudspace_id, + ) def wait(self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False) -> None: import time @@ -560,14 +558,12 @@ def rank(self) -> Optional[int]: @property def is_multi_machine(self) -> bool: """Whether this object represents a multi-machine parent job.""" - return self._resource_kind == "multi" + return self._num_machines > 1 @property def num_machines(self) -> int: """The number of machines allocated to this job.""" - if not self.is_multi_machine: - return 1 - return self._job_api.get_num_machines(self._guaranteed_job) + return self._num_machines @property def machines(self) -> Tuple["Job", ...]: @@ -581,8 +577,8 @@ def machines(self) -> Tuple["Job", ...]: ) machines = [] for subjob in subjobs: - job = Job(name=subjob.name, teamspace=self.teamspace, _fetch_job=False) - job._job = subjob + job = Job(name=subjob.name, teamspace=self.teamspace, _fetch_job=False, _num_machines=1) + job._attach_job(subjob) machines.append(job) return tuple(machines) @@ -911,21 +907,24 @@ def command(self) -> str: def _update_internal_job(self) -> None: if getattr(self, "_job", None) is None: if self.is_multi_machine: - self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) + self._attach_job(self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id)) return from lightning_sdk.lightning_cloud.openapi.rest import ApiException try: - self._job = self._standalone_job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) + self._attach_job( + self._standalone_job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) + ) except ApiException as ex: if ex.status != 404: raise - self._set_resource_kind("multi") - self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) + # Switch to the multi-machine API, then sync the real machine count from the payload. + self._num_machines = 2 + self._attach_job(self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id)) return - self._job = self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id) + self._attach_job(self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id)) @property def name(self) -> str: diff --git a/python/lightning_sdk/mmt.py b/python/lightning_sdk/mmt.py index 31df6b67..c92d02ff 100644 --- a/python/lightning_sdk/mmt.py +++ b/python/lightning_sdk/mmt.py @@ -70,6 +70,7 @@ def __init__( user: Union[str, "User", None] = None, *, _fetch_job: bool = True, + _num_machines: int = 2, ) -> None: try: super().__init__( @@ -78,7 +79,8 @@ def __init__( org=org, user=user, _fetch_job=_fetch_job, - _resource_kind="multi", + # Default 2 forces the multi-machine API for lookup; real count is synced after fetch/attach. + _num_machines=_num_machines, ) except ValueError as ex: # Job.__init__ raises "Job {name} does not exist…" on 404; keep the MMT-specific diff --git a/python/lightning_sdk/teamspace.py b/python/lightning_sdk/teamspace.py index 75ff34c0..668c3391 100644 --- a/python/lightning_sdk/teamspace.py +++ b/python/lightning_sdk/teamspace.py @@ -294,11 +294,11 @@ def jobs(self) -> Tuple["Job", ...]: for j2 in self._teamspace_api.list_jobs(teamspace_id=self.id): # _fetch_job = False to prevent refetching on init since we already got it job = Job(name=j2.name, teamspace=self, _fetch_job=False) - job._job = j2 + job._attach_job(j2) jobs.append(job) for m2 in self._teamspace_api.list_mmts(teamspace_id=self.id): - job = Job(name=m2.name, teamspace=self, _fetch_job=False, _resource_kind="multi") - job._job = m2 + job = Job(name=m2.name, teamspace=self, _fetch_job=False) + job._attach_job(m2) jobs.append(job) return tuple(jobs) @@ -321,7 +321,7 @@ def multi_machine_jobs(self) -> Tuple["MMT", ...]: for m2 in self._teamspace_api.list_mmts(teamspace_id=self.id): # _fetch_job = False to prevent refetching on init since we already got it mmt = MMT(name=m2.name, teamspace=self, _fetch_job=False) - mmt._job = m2 + mmt._attach_job(m2) mmts.append(mmt) return tuple(mmts) diff --git a/python/tests/core/test_job.py b/python/tests/core/test_job.py index 4e074d30..92abdb82 100644 --- a/python/tests/core/test_job.py +++ b/python/tests/core/test_job.py @@ -159,6 +159,7 @@ def test_submit_job_v2_image(internal_studio_init_mocker, machine, command, env, reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) @@ -245,6 +246,7 @@ def test_submit_job_v2_studio(internal_studio_init_mocker, machine, env, interru reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) diff --git a/python/tests/core/test_mmt.py b/python/tests/core/test_mmt.py index 6d60d26d..c7ad87e2 100644 --- a/python/tests/core/test_mmt.py +++ b/python/tests/core/test_mmt.py @@ -57,6 +57,7 @@ def test_submit_mmt_v2_image(internal_studio_init_mocker, machine, command, env, max_runtime=None, reuse_snapshot=True, placement_group_id=None, + scratch_disks=None, ) @@ -172,6 +173,7 @@ def test_submit_mmt_v2_studio(internal_studio_init_mocker, machine, env, interru max_runtime=None, reuse_snapshot=True, placement_group_id=None, + scratch_disks=None, ) @@ -473,7 +475,7 @@ def test_mmtv2_delete(mmt_api_get_job_by_name_mocker, internal_studio_init_mocke job.delete() - delete_job_mock.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001") + delete_job_mock.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001", cloudspace_id=None) @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) diff --git a/python/tests/core/test_unified_job.py b/python/tests/core/test_unified_job.py index ccda2162..6f13a13c 100644 --- a/python/tests/core/test_unified_job.py +++ b/python/tests/core/test_unified_job.py @@ -52,7 +52,7 @@ def test_job_lookup_prefers_standalone_on_name_collision() -> None: assert not job.is_multi_machine assert job.num_machines == 1 assert job.machines == (job,) - multi_api.assert_not_called() + multi_api.return_value.get_job_by_name.assert_not_called() def test_job_run_routes_multi_machine_submission() -> None: