Skip to content

Commit 9a771e7

Browse files
authored
feat: update functionality for job.log with name argument (#146)
* feat(event): add/update event log with name Can use the event's ID as its name, named events reset when a checkpoint is created * test(event): tests for updating logs * test(event): cleanup * fix(event): typo for python 3.8 * fix(event): typo for python 3.8 * fix(event): review fixes * fix(event): revert to name/id behaviour * docs(event): better name docstring * refactor(event): separate name and ID arguments No longer use ID as name to reduce ambiguity * docs(event): fix docstring typo
1 parent ddbf2bd commit 9a771e7

File tree

2 files changed

+143
-5
lines changed

2 files changed

+143
-5
lines changed

cryosparc/controllers/job.py

Lines changed: 61 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,11 @@
5757
Input and output result groups may only contain, letters, numbers and underscores.
5858
"""
5959

60+
LogLevel = Literal["text", "warning", "error"]
61+
"""
62+
Severity level for job event logs.
63+
"""
64+
6065

6166
class JobController(Controller[Job]):
6267
"""
@@ -108,6 +113,15 @@ class JobController(Controller[Job]):
108113
Project unique ID, e.g., "P3"
109114
"""
110115

116+
_events: Dict[str, str]
117+
"""
118+
Named event logs. Key can be user-provided name in log() method, or ID if
119+
name not provided. If both name and ID are used, can have two keys with
120+
the same value.
121+
122+
:meta private:
123+
"""
124+
111125
def __init__(self, cs: "CryoSPARC", job: Union[Tuple[str, str], Job]) -> None:
112126
self.cs = cs
113127
if isinstance(job, tuple):
@@ -117,6 +131,7 @@ def __init__(self, cs: "CryoSPARC", job: Union[Tuple[str, str], Job]) -> None:
117131
self.project_uid = job.project_uid
118132
self.uid = job.uid
119133
self.model = job
134+
self._events = {}
120135

121136
@property
122137
def type(self) -> str:
@@ -526,24 +541,65 @@ def load_output(self, name: str, slots: LoadableSlots = "all", version: Union[in
526541
"""
527542
return self.cs.api.jobs.load_output(self.project_uid, self.uid, name, slots=slots, version=version)
528543

529-
def log(self, text: str, level: Literal["text", "warning", "error"] = "text"):
544+
@overload
545+
def log(self, text: str, *, level: LogLevel = ...) -> str: ...
546+
@overload
547+
def log(self, text: str, *, level: LogLevel = ..., name: str) -> str: ...
548+
@overload
549+
def log(self, text: str, *, level: LogLevel = ..., id: str) -> str: ...
550+
def log(self, text: str, *, level: LogLevel = "text", name: Optional[str] = None, id: Optional[str] = None) -> str:
530551
"""
531-
Append to a job's event log.
552+
Append to a job's event log. Update an existing log by providing a name
553+
or ID.
532554
533555
Args:
534556
text (str): Text to log
535557
level (str, optional): Log level ("text", "warning" or "error").
536558
Defaults to "text".
559+
name (str, optional): Event name. If called multiple times with the
560+
same name, updates that event instead of creating a new one.
561+
Named events are reset when logging a checkpoint. Cannot be
562+
provided with id. Defaults to None.
563+
id (str, optional): Update a previously-created event log by its ID.
564+
Cannot be provided with name. Defaults to None.
565+
566+
Example:
567+
568+
Log a warning message to the job log.
569+
>>> job.log("This is a warning", level="warning")
570+
571+
Show a live progress bar in the job log.
572+
>>> for pct in range(1, 10):
573+
... # example log: "Progress: [#####-----] 50%"
574+
... job.log(f"Progress: [{'#' * pct}{'-' * (10 - pct)}] {pct * 10}%", name="progress")
575+
... sleep(1)
576+
...
577+
>>> job.log("Done!")
578+
579+
Update an existing log event by ID.
580+
>>> event_id = job.log("Starting job processing...")
581+
>>> # do some processing...
582+
>>> job.log("Finished processing", id=event_id)
537583
538584
Returns:
539585
str: Created log event ID
540586
"""
541-
event = self.cs.api.jobs.add_event_log(self.project_uid, self.uid, text, type=level)
587+
existing_id = id
588+
if name and name in self._events:
589+
existing_id = self._events[name]
590+
591+
if existing_id:
592+
event = self.cs.api.jobs.update_event_log(self.project_uid, self.uid, existing_id, text, type=level)
593+
else:
594+
event = self.cs.api.jobs.add_event_log(self.project_uid, self.uid, text, type=level)
595+
596+
if name:
597+
self._events[name] = event.id
542598
return event.id
543599

544600
def log_checkpoint(self, meta: dict = {}):
545601
"""
546-
Append a checkpoint to the job's event log.
602+
Append a checkpoint to the job's event log. Also resets named events.
547603
548604
Args:
549605
meta (dict, optional): Additional meta information. Defaults to {}.
@@ -552,6 +608,7 @@ def log_checkpoint(self, meta: dict = {}):
552608
str: Created checkpoint event ID
553609
"""
554610
event = self.cs.api.jobs.add_checkpoint(self.project_uid, self.uid, meta)
611+
self._events = {}
555612
return event.id
556613

557614
def log_plot(

tests/controllers/test_job.py

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,6 @@ def test_job_subprocess_io(job: JobController):
8080
[sys.executable, "-c", 'import sys; print("hello"); print("error", file=sys.stderr); print("world")']
8181
)
8282

83-
assert len(mock_log_endpoint.mock_calls) == 7 # includes some prelude/divider calls
8483
mock_log_endpoint.assert_has_calls(
8584
[
8685
mock.call(job.project_uid, job.uid, "hello", type="text"),
@@ -185,3 +184,85 @@ def test_external_job_output(mock_external_job_with_saved_output: ExternalJobCon
185184
def test_invalid_external_job_output(external_job):
186185
with pytest.raises(ValueError, match="Invalid output name"):
187186
external_job.add_output("particle", name="particles/1", slots=["blob", "ctf"])
187+
188+
189+
@pytest.fixture
190+
def mock_log_event():
191+
return mock.MagicMock(id="event_123")
192+
193+
194+
@pytest.fixture
195+
def mock_checkpoint_event():
196+
return mock.MagicMock(id="checkpoint_456")
197+
198+
199+
def test_log(job: JobController, mock_log_event):
200+
assert isinstance(mock_add_endpoint := APIClient.jobs.add_event_log, mock.Mock)
201+
mock_add_endpoint.return_value = mock_log_event
202+
203+
result = job.log("Test message without name")
204+
205+
mock_add_endpoint.assert_called_once_with(job.project_uid, job.uid, "Test message without name", type="text")
206+
assert result == mock_log_event.id
207+
208+
209+
def test_log_with_name_create_and_update(job: JobController, mock_log_event):
210+
assert isinstance(mock_add_endpoint := APIClient.jobs.add_event_log, mock.Mock)
211+
assert isinstance(mock_update_endpoint := APIClient.jobs.update_event_log, mock.Mock)
212+
mock_add_endpoint.return_value = mock_log_event
213+
mock_update_endpoint.return_value = mock_log_event
214+
215+
# First call with name - should create
216+
event_id = job.log("First message", name="progress")
217+
218+
mock_add_endpoint.assert_called_once_with(job.project_uid, job.uid, "First message", type="text")
219+
assert event_id == "event_123"
220+
221+
# Second call with same name - should update
222+
event_id = job.log("Updated message", level="warning", name="progress")
223+
224+
mock_update_endpoint.assert_called_once_with(
225+
job.project_uid, job.uid, mock_log_event.id, "Updated message", type="warning"
226+
)
227+
assert event_id == "event_123"
228+
229+
230+
def test_log_with_returned_event_id(job: JobController, mock_log_event):
231+
assert isinstance(mock_add_endpoint := APIClient.jobs.add_event_log, mock.Mock)
232+
assert isinstance(mock_update_endpoint := APIClient.jobs.update_event_log, mock.Mock)
233+
mock_add_endpoint.return_value = mock_log_event
234+
mock_update_endpoint.return_value = mock_log_event
235+
236+
# First call without ID - returns event ID
237+
event_id = job.log("Initial message")
238+
assert event_id == mock_log_event.id
239+
240+
# Second call using the returned event ID - should update
241+
result = job.log("Updated with event ID", id=event_id)
242+
243+
mock_update_endpoint.assert_called_once_with(
244+
job.project_uid, job.uid, mock_log_event.id, "Updated with event ID", type="text"
245+
)
246+
assert result == event_id
247+
248+
249+
def test_log_after_checkpoint_creates_new(job: JobController, mock_log_event, mock_checkpoint_event):
250+
assert isinstance(mock_add_endpoint := APIClient.jobs.add_event_log, mock.Mock)
251+
assert isinstance(mock_update_endpoint := APIClient.jobs.update_event_log, mock.Mock)
252+
assert isinstance(mock_checkpoint_endpoint := APIClient.jobs.add_checkpoint, mock.Mock)
253+
254+
mock_add_endpoint.return_value = mock_log_event
255+
mock_update_endpoint.return_value = mock_log_event
256+
mock_checkpoint_endpoint.return_value = mock_checkpoint_event
257+
258+
job.log("Before checkpoint", name="status")
259+
260+
checkpoint_id = job.log_checkpoint()
261+
mock_checkpoint_endpoint.assert_called_once_with(job.project_uid, job.uid, {})
262+
assert checkpoint_id == mock_checkpoint_event.id
263+
264+
# Log again with same name - should create new
265+
mock_add_endpoint.reset_mock() # Reset to track the second call
266+
result = job.log("After checkpoint", name="status")
267+
mock_add_endpoint.assert_called_once_with(job.project_uid, job.uid, "After checkpoint", type="text")
268+
assert result == "event_123"

0 commit comments

Comments
 (0)