Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions go/job.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ type Job struct {
mmtID string
publicIP string
totalCost float32
startedAt time.Time
stoppedAt time.Time

artifactsSource string
artifactsDestination string
Expand Down Expand Up @@ -270,6 +272,24 @@ func (j *Job) TotalCost() float32 {
return j.totalCost
}

// StartedAt returns when the job started running, or the zero time if it
// hasn't started yet.
func (j *Job) StartedAt() time.Time {
if j == nil {
return time.Time{}
}
return j.startedAt
}

// StoppedAt returns when the job stopped running, or the zero time if it
// hasn't stopped yet.
func (j *Job) StoppedAt() time.Time {
if j == nil {
return time.Time{}
}
return j.stoppedAt
}

// ArtifactsSource returns the local artifacts source path.
func (j *Job) ArtifactsSource() string {
if j == nil {
Expand Down Expand Up @@ -705,6 +725,8 @@ func jobFromModel(model *models.V1Job, opts jobOptions) *Job {
mmtID: model.MultiMachineJobID,
publicIP: model.PublicIPAddress,
totalCost: model.TotalCost,
startedAt: time.Time(model.StartedAt),
stoppedAt: time.Time(model.StoppedAt),
}
if model.Spec != nil {
result.machine = model.Spec.InstanceName
Expand Down
24 changes: 24 additions & 0 deletions go/job_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,37 @@ import (
"net/http/httptest"
"strings"
"testing"
"time"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

lit "github.com/lightning-ai/sdk/go"
)

func TestJobExposesStartAndStopTimes(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "job-1",
"name": "train",
"projectId": "project-1",
"state": "running",
"startedAt": "2026-08-01T12:00:00Z",
})
}))
defer server.Close()
t.Setenv("LIGHTNING_CLOUD_URL", server.URL)

existing, err := lit.GetJob("train", lit.JobOptions{Teamspace: mustTeamspace(t, "project-1", "default", "alice")})
require.NoErrorf(t, err,
"GetJob returned error")
assert.Truef(t, existing.StartedAt().Equal(time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)),
"StartedAt = %v, want 2026-08-01T12:00:00Z", existing.StartedAt())
assert.Truef(t, existing.StoppedAt().IsZero(),
"StoppedAt = %v, want zero time", existing.StoppedAt())
}

func TestJobGetWithIDUsesSimpleStruct(t *testing.T) {
j, err := lit.GetJob("train", lit.JobOptions{ID: "job-1", Teamspace: mustTeamspace(t, "project-1", "")})
require.NoErrorf(t, err,
Expand Down
24 changes: 24 additions & 0 deletions go/mmt.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@ type MMT struct {
image string
studioID string
totalCost float32
startedAt time.Time
stoppedAt time.Time
}

// MachineDict is the JSON-friendly public representation of one MMT machine.
Expand Down Expand Up @@ -249,6 +251,24 @@ func (m *MMT) TotalCost() float32 {
return m.totalCost
}

// StartedAt returns when the MMT's machines started running, or the zero time
// if they haven't started yet.
func (m *MMT) StartedAt() time.Time {
if m == nil {
return time.Time{}
}
return m.startedAt
}

// StoppedAt returns when the MMT stopped running, or the zero time if it
// hasn't stopped yet.
func (m *MMT) StoppedAt() time.Time {
if m == nil {
return time.Time{}
}
return m.stoppedAt
}

// GetMMT returns an existing MMT by name or ID.
func GetMMT(name string, opts ...MMTOptions) (*MMT, error) {
resolved := applyMMTOptions(opts...)
Expand Down Expand Up @@ -632,6 +652,10 @@ func mmtFromModel(model *models.V1MultiMachineJob, opts mmtOptions) *MMT {
if model.State != nil {
result.status = string(*model.State)
}
if model.Status != nil {
result.startedAt = time.Time(model.Status.StartedAt)
result.stoppedAt = time.Time(model.Status.StoppedAt)
}
if model.Spec != nil {
result.machine = model.Spec.InstanceName
result.command = model.Spec.Command
Expand Down
28 changes: 28 additions & 0 deletions go/mmt_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"net/http/httptest"
"strings"
"testing"
"time"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
Expand Down Expand Up @@ -337,6 +338,33 @@ func TestMMTGetMapsTopLevelStudioID(t *testing.T) {

}

func TestMMTExposesStartAndStopTimes(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "mmt-1",
"name": "dist-train",
"projectId": "project-1",
"machines": 4,
"state": "stopped",
"status": map[string]any{
"startedAt": "2026-08-02T12:00:00Z",
"stoppedAt": "2026-08-02T13:00:00Z",
},
})
}))
defer server.Close()
t.Setenv("LIGHTNING_CLOUD_URL", server.URL)

existing, err := lit.GetMMT("dist-train", lit.MMTOptions{Teamspace: mustTeamspace(t, "project-1", "")})
require.NoErrorf(t, err,
"GetMMT returned error")
assert.Truef(t, existing.StartedAt().Equal(time.Date(2026, 8, 2, 12, 0, 0, 0, time.UTC)),
"StartedAt = %v, want 2026-08-02T12:00:00Z", existing.StartedAt())
assert.Truef(t, existing.StoppedAt().Equal(time.Date(2026, 8, 2, 13, 0, 0, 0, time.UTC)),
"StoppedAt = %v, want 2026-08-02T13:00:00Z", existing.StoppedAt())
}

func TestMMTRunMapsAdvancedV2Options(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
Expand Down
41 changes: 37 additions & 4 deletions python/lightning_sdk/cli/job/list.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Job list command."""

from contextlib import suppress
from datetime import datetime
from typing import Optional

import rich_click as click
Expand Down Expand Up @@ -34,7 +35,8 @@
"--sort_by",
default=None,
type=click.Choice(
["name", "teamspace", "status", "studio", "machine", "image", "cloud-account"], case_sensitive=False
["name", "teamspace", "status", "studio", "machine", "image", "cloud-account", "started", "stopped"],
case_sensitive=False,
),
help="the attribute to sort the jobs by.",
)
Expand Down Expand Up @@ -69,21 +71,46 @@ def list_jobs(
"studio": job.studio_name,
"image": job.image,
"status": str(job.status) if job.status is not None else None,
"started_at": getattr(job, "started_at", None),
"stopped_at": getattr(job, "stopped_at", None),
"machine": str(job.machine),
"num_machines": getattr(job, "num_machines", 1),
"total_cost": round(job.total_cost, 3),
"_cloud_account": str(getattr(job, "cloud_account", "") or ""),
}
)

sort_key = "_cloud_account" if sort_by == "cloud-account" else sort_by or "name"
sort_by = sort_by or "name"
sort_key = {"cloud-account": "_cloud_account", "started": "started_at", "stopped": "stopped_at"}.get(
sort_by, sort_by
)
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])
echo_json(
[
{
key: value.isoformat() if isinstance(value, datetime) else 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"):
for column in (
"Name",
"Teamspace",
"Studio",
"Image",
"Status",
"Started",
"Stopped",
"Machine",
"Num Machines",
"Total Cost",
):
table.add_column(column)
for row in rows:
table.add_row(
Expand All @@ -92,8 +119,14 @@ def list_jobs(
str(row["studio"] or ""),
str(row["image"] or ""),
str(row["status"] or ""),
_format_timestamp(row["started_at"]),
_format_timestamp(row["stopped_at"]),
str(row["machine"] or ""),
str(row["num_machines"]),
f"{row['total_cost']:.3f}",
)
Console().print(table)


def _format_timestamp(value: object) -> str:
return value.strftime("%Y-%m-%d %H:%M") if isinstance(value, datetime) else ""
20 changes: 20 additions & 0 deletions python/lightning_sdk/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
)

if TYPE_CHECKING:
from datetime import datetime

from lightning_sdk.machine import CloudProvider, Machine
from lightning_sdk.organization import Organization
from lightning_sdk.studio import Studio
Expand Down Expand Up @@ -534,6 +536,24 @@ def status(self) -> Status:
f"Job {self._name} does not exist in Teamspace {self.teamspace.name}. Did you delete it?"
) from None

@property
def started_at(self) -> Optional["datetime"]:
"""When the job started running, or ``None`` if it hasn't started yet."""
job = self._latest_job
if self.is_multi_machine:
status = getattr(job, "status", None)
return status.started_at if status is not None else None
return getattr(job, "started_at", None)

@property
def stopped_at(self) -> Optional["datetime"]:
"""When the job stopped running, or ``None`` if it hasn't stopped yet."""
job = self._latest_job
if self.is_multi_machine:
status = getattr(job, "status", None)
return status.stopped_at if status is not None else None
return getattr(job, "stopped_at", None)

@property
def machine(self) -> Union["Machine", str]:
return self._job_api._get_job_machine_from_spec(
Expand Down
31 changes: 31 additions & 0 deletions python/tests/cli/job/test_list.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import json
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import patch

Expand Down Expand Up @@ -37,6 +38,8 @@ def _teamspace_with_jobs() -> SimpleNamespace:
image="ubuntu",
status="Running",
machine="CPU",
started_at=datetime(2026, 8, 1, 12, 0, tzinfo=timezone.utc),
stopped_at=datetime(2026, 8, 1, 13, 0, tzinfo=timezone.utc),
total_cost=1.0,
)
multi = SimpleNamespace(
Expand All @@ -47,6 +50,8 @@ def _teamspace_with_jobs() -> SimpleNamespace:
status="Running",
machine="CPU",
num_machines=4,
started_at=datetime(2026, 8, 2, 12, 0, tzinfo=timezone.utc),
stopped_at=None,
total_cost=4.0,
)
teamspace.jobs = [single, multi]
Expand All @@ -69,6 +74,32 @@ def test_job_list_includes_single_and_multi_machine_jobs() -> None:
]


@mock_command_logging
def test_job_list_includes_and_sorts_by_timestamps() -> None:
teamspace = _teamspace_with_jobs()

with patch("lightning_sdk.cli.job.list.resolve_teamspace", return_value=teamspace):
result = CliRunner().invoke(list_jobs, ["--sort-by", "started", "--json"])

assert result.exit_code == 0, result.output
rows = json.loads(result.output)
assert [row["name"] for row in rows] == ["single", "distributed"]
assert rows[0]["started_at"] == "2026-08-01T12:00:00+00:00"
assert rows[0]["stopped_at"] == "2026-08-01T13:00:00+00:00"
assert rows[1]["stopped_at"] is None


@mock_command_logging
def test_job_list_sort_by_stopped_puts_unfinished_first() -> None:
teamspace = _teamspace_with_jobs()

with patch("lightning_sdk.cli.job.list.resolve_teamspace", return_value=teamspace):
result = CliRunner().invoke(list_jobs, ["--sort-by", "stopped", "--json"])

assert result.exit_code == 0, result.output
assert [row["name"] for row in json.loads(result.output)] == ["distributed", "single"]


@mock_command_logging
def test_job_list_sort_by_cloud_account_without_attribute() -> None:
teamspace = _teamspace_with_jobs()
Expand Down
42 changes: 42 additions & 0 deletions python/tests/core/test_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
JobsServiceUpdateJobBody,
V1Job,
V1JobSpec,
V1MultiMachineJob,
V1MultiMachineJobStatus,
)
from lightning_sdk.lightning_cloud.openapi.rest import ApiException
from lightning_sdk.machine import Machine
Expand Down Expand Up @@ -256,6 +258,46 @@ def test_job_exposes_private_provisioning_metadata(internal_studio_init_mocker):
assert job.rank == 3


@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock())
def test_job_exposes_start_and_stop_times(internal_studio_init_mocker):
teamspace = Teamspace("ts-abc", org="org-abc")
job = Job("test-job", teamspace, _fetch_job=False)
job._prevent_refetch_latest = True
job._attach_job(
V1Job(
id="job-123",
name="test-job",
started_at=datetime(2026, 8, 1, 12, 0, tzinfo=timezone.utc),
spec=V1JobSpec(),
)
)

assert job.started_at == datetime(2026, 8, 1, 12, 0, tzinfo=timezone.utc)
assert job.stopped_at is None


@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock())
def test_multi_machine_job_exposes_start_and_stop_times(internal_studio_init_mocker):
teamspace = Teamspace("ts-abc", org="org-abc")
job = Job("test-mmt", teamspace, _fetch_job=False)
job._prevent_refetch_latest = True
job._attach_job(
V1MultiMachineJob(
id="mmt-123",
name="test-mmt",
machines=4,
status=V1MultiMachineJobStatus(
started_at=datetime(2026, 8, 2, 12, 0, tzinfo=timezone.utc),
stopped_at=datetime(2026, 8, 2, 13, 0, tzinfo=timezone.utc),
),
)
)

assert job.is_multi_machine
assert job.started_at == datetime(2026, 8, 2, 12, 0, tzinfo=timezone.utc)
assert job.stopped_at == datetime(2026, 8, 2, 13, 0, tzinfo=timezone.utc)


@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock())
def test_job_id_is_none_without_backing_model_or_api_request(internal_studio_init_mocker):
teamspace = Teamspace("org-abc/ts-abc")
Expand Down
Loading