UncountablePythonSDK 0.0.126__py3-none-any.whl → 0.0.128__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of UncountablePythonSDK might be problematic. Click here for more details.

Files changed (49) hide show
  1. examples/integration-server/jobs/materials_auto/example_instrument.py +67 -38
  2. examples/integration-server/jobs/materials_auto/example_parse.py +87 -0
  3. examples/integration-server/jobs/materials_auto/example_predictions.py +61 -0
  4. examples/integration-server/jobs/materials_auto/example_runsheet_wh.py +3 -2
  5. examples/integration-server/jobs/materials_auto/profile.yaml +18 -0
  6. examples/integration-server/pyproject.toml +3 -3
  7. pkgs/type_spec/builder.py +19 -9
  8. pkgs/type_spec/emit_typescript.py +2 -2
  9. pkgs/type_spec/type_info/emit_type_info.py +14 -1
  10. pkgs/type_spec/value_spec/__main__.py +2 -2
  11. uncountable/integration/cli.py +29 -1
  12. uncountable/integration/executors/executors.py +1 -2
  13. uncountable/integration/executors/generic_upload_executor.py +1 -1
  14. uncountable/integration/job.py +3 -3
  15. uncountable/integration/queue_runner/command_server/__init__.py +4 -0
  16. uncountable/integration/queue_runner/command_server/command_client.py +39 -0
  17. uncountable/integration/queue_runner/command_server/command_server.py +37 -0
  18. uncountable/integration/queue_runner/command_server/protocol/command_server.proto +18 -0
  19. uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.py +21 -13
  20. uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.pyi +28 -1
  21. uncountable/integration/queue_runner/command_server/protocol/command_server_pb2_grpc.py +90 -0
  22. uncountable/integration/queue_runner/command_server/types.py +24 -1
  23. uncountable/integration/queue_runner/datastore/datastore_sqlite.py +107 -8
  24. uncountable/integration/queue_runner/datastore/model.py +8 -1
  25. uncountable/integration/queue_runner/job_scheduler.py +42 -2
  26. uncountable/integration/queue_runner/worker.py +1 -1
  27. uncountable/integration/server.py +36 -6
  28. uncountable/integration/telemetry.py +41 -7
  29. uncountable/types/__init__.py +4 -0
  30. uncountable/types/api/integrations/register_sockets_token.py +41 -0
  31. uncountable/types/api/recipes/edit_recipe_inputs.py +1 -1
  32. uncountable/types/api/recipes/get_recipes_data.py +16 -0
  33. uncountable/types/api/recipes/lock_recipes.py +2 -1
  34. uncountable/types/api/recipes/set_recipe_total.py +59 -0
  35. uncountable/types/api/recipes/unlock_recipes.py +2 -1
  36. uncountable/types/api/uploader/complete_async_parse.py +4 -0
  37. uncountable/types/async_batch_processor.py +124 -0
  38. uncountable/types/async_batch_t.py +2 -0
  39. uncountable/types/client_base.py +57 -1
  40. uncountable/types/entity_t.py +1 -1
  41. uncountable/types/queued_job.py +1 -0
  42. uncountable/types/queued_job_t.py +9 -0
  43. uncountable/types/sockets.py +9 -0
  44. uncountable/types/sockets_t.py +99 -0
  45. uncountable/types/uploader_t.py +3 -2
  46. {uncountablepythonsdk-0.0.126.dist-info → uncountablepythonsdk-0.0.128.dist-info}/METADATA +1 -1
  47. {uncountablepythonsdk-0.0.126.dist-info → uncountablepythonsdk-0.0.128.dist-info}/RECORD +49 -45
  48. {uncountablepythonsdk-0.0.126.dist-info → uncountablepythonsdk-0.0.128.dist-info}/WHEEL +0 -0
  49. {uncountablepythonsdk-0.0.126.dist-info → uncountablepythonsdk-0.0.128.dist-info}/top_level.txt +0 -0
@@ -13,11 +13,19 @@ from uncountable.integration.queue_runner.command_server.protocol.command_server
13
13
  EnqueueJobResult,
14
14
  ListQueuedJobsRequest,
15
15
  ListQueuedJobsResult,
16
+ RetryJobRequest,
17
+ RetryJobResult,
18
+ VaccuumQueuedJobsRequest,
19
+ VaccuumQueuedJobsResult,
16
20
  )
17
21
  from uncountable.integration.queue_runner.command_server.types import (
18
22
  CommandEnqueueJob,
19
23
  CommandEnqueueJobResponse,
20
24
  CommandQueue,
25
+ CommandRetryJob,
26
+ CommandRetryJobResponse,
27
+ CommandVaccuumQueuedJobs,
28
+ CommandVaccuumQueuedJobsResponse,
21
29
  )
22
30
  from uncountable.integration.queue_runner.datastore import DatastoreSqlite
23
31
  from uncountable.types import queued_job_t
@@ -54,6 +62,23 @@ async def serve(command_queue: CommandQueue, datastore: DatastoreSqlite) -> None
54
62
  )
55
63
  return result
56
64
 
65
+ async def RetryJob(
66
+ self, request: RetryJobRequest, context: aio.ServicerContext
67
+ ) -> RetryJobResult:
68
+ response_queue: asyncio.Queue[CommandRetryJobResponse] = asyncio.Queue()
69
+ await command_queue.put(
70
+ CommandRetryJob(
71
+ queued_job_uuid=request.uuid, response_queue=response_queue
72
+ )
73
+ )
74
+ response = await response_queue.get()
75
+ if response.queued_job_uuid is not None:
76
+ return RetryJobResult(
77
+ successfully_queued=True, queued_job_uuid=response.queued_job_uuid
78
+ )
79
+ else:
80
+ return RetryJobResult(successfully_queued=False, queued_job_uuid="")
81
+
57
82
  async def CheckHealth(
58
83
  self, request: CheckHealthRequest, context: aio.ServicerContext
59
84
  ) -> CheckHealthResult:
@@ -90,10 +115,22 @@ async def serve(command_queue: CommandQueue, datastore: DatastoreSqlite) -> None
90
115
  job_ref_name=item.job_ref_name,
91
116
  num_attempts=item.num_attempts,
92
117
  submitted_at=proto_timestamp,
118
+ status=item.status,
93
119
  )
94
120
  )
95
121
  return ListQueuedJobsResult(queued_jobs=response_list)
96
122
 
123
+ async def VaccuumQueuedJobs(
124
+ self, request: VaccuumQueuedJobsRequest, context: aio.ServicerContext
125
+ ) -> VaccuumQueuedJobsResult:
126
+ response_queue: asyncio.Queue[CommandVaccuumQueuedJobsResponse] = (
127
+ asyncio.Queue()
128
+ )
129
+ await command_queue.put(
130
+ CommandVaccuumQueuedJobs(response_queue=response_queue)
131
+ )
132
+ return VaccuumQueuedJobsResult()
133
+
97
134
  add_CommandServerServicer_to_server(CommandServerHandler(), server)
98
135
 
99
136
  listen_addr = f"[::]:{get_local_admin_server_port()}"
@@ -3,8 +3,10 @@ import "google/protobuf/timestamp.proto";
3
3
 
4
4
  service CommandServer {
5
5
  rpc EnqueueJob(EnqueueJobRequest) returns (EnqueueJobResult) {}
6
+ rpc RetryJob(RetryJobRequest) returns (RetryJobResult) {}
6
7
  rpc CheckHealth(CheckHealthRequest) returns (CheckHealthResult) {}
7
8
  rpc ListQueuedJobs(ListQueuedJobsRequest) returns (ListQueuedJobsResult) {}
9
+ rpc VaccuumQueuedJobs(VaccuumQueuedJobsRequest) returns (VaccuumQueuedJobsResult) {}
8
10
  }
9
11
 
10
12
  message EnqueueJobRequest {
@@ -17,6 +19,21 @@ message EnqueueJobResult {
17
19
  string queued_job_uuid = 2;
18
20
  }
19
21
 
22
+ message RetryJobRequest {
23
+ string uuid = 1;
24
+ }
25
+
26
+ message RetryJobResult {
27
+ bool successfully_queued = 1;
28
+ string queued_job_uuid = 2;
29
+ }
30
+
31
+ message VaccuumQueuedJobsRequest {
32
+ }
33
+
34
+ message VaccuumQueuedJobsResult {
35
+ }
36
+
20
37
  message CheckHealthRequest {}
21
38
 
22
39
  message CheckHealthResult {
@@ -34,6 +51,7 @@ message ListQueuedJobsResult {
34
51
  string job_ref_name = 2;
35
52
  int64 num_attempts = 3;
36
53
  google.protobuf.Timestamp submitted_at = 4;
54
+ string status = 5;
37
55
  }
38
56
 
39
57
  repeated ListQueuedJobsResultItem queued_jobs = 1;
@@ -18,7 +18,7 @@ from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__
18
18
 
19
19
 
20
20
  DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
21
- b'\nQuncountable/integration/queue_runner/command_server/protocol/command_server.proto\x1a\x1fgoogle/protobuf/timestamp.proto"E\n\x11\x45nqueueJobRequest\x12\x14\n\x0cjob_ref_name\x18\x01 \x01(\t\x12\x1a\n\x12serialized_payload\x18\x02 \x01(\t"H\n\x10\x45nqueueJobResult\x12\x1b\n\x13successfully_queued\x18\x01 \x01(\x08\x12\x17\n\x0fqueued_job_uuid\x18\x02 \x01(\t"\x14\n\x12\x43heckHealthRequest"$\n\x11\x43heckHealthResult\x12\x0f\n\x07success\x18\x01 \x01(\x08"6\n\x15ListQueuedJobsRequest\x12\x0e\n\x06offset\x18\x01 \x01(\r\x12\r\n\x05limit\x18\x02 \x01(\r"\xe4\x01\n\x14ListQueuedJobsResult\x12\x43\n\x0bqueued_jobs\x18\x01 \x03(\x0b\x32..ListQueuedJobsResult.ListQueuedJobsResultItem\x1a\x86\x01\n\x18ListQueuedJobsResultItem\x12\x0c\n\x04uuid\x18\x01 \x01(\t\x12\x14\n\x0cjob_ref_name\x18\x02 \x01(\t\x12\x14\n\x0cnum_attempts\x18\x03 \x01(\x03\x12\x30\n\x0csubmitted_at\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.Timestamp2\xc3\x01\n\rCommandServer\x12\x35\n\nEnqueueJob\x12\x12.EnqueueJobRequest\x1a\x11.EnqueueJobResult"\x00\x12\x38\n\x0b\x43heckHealth\x12\x13.CheckHealthRequest\x1a\x12.CheckHealthResult"\x00\x12\x41\n\x0eListQueuedJobs\x12\x16.ListQueuedJobsRequest\x1a\x15.ListQueuedJobsResult"\x00\x62\x06proto3'
21
+ b'\nQuncountable/integration/queue_runner/command_server/protocol/command_server.proto\x1a\x1fgoogle/protobuf/timestamp.proto"E\n\x11\x45nqueueJobRequest\x12\x14\n\x0cjob_ref_name\x18\x01 \x01(\t\x12\x1a\n\x12serialized_payload\x18\x02 \x01(\t"H\n\x10\x45nqueueJobResult\x12\x1b\n\x13successfully_queued\x18\x01 \x01(\x08\x12\x17\n\x0fqueued_job_uuid\x18\x02 \x01(\t"\x1f\n\x0fRetryJobRequest\x12\x0c\n\x04uuid\x18\x01 \x01(\t"F\n\x0eRetryJobResult\x12\x1b\n\x13successfully_queued\x18\x01 \x01(\x08\x12\x17\n\x0fqueued_job_uuid\x18\x02 \x01(\t"\x1a\n\x18VaccuumQueuedJobsRequest"\x19\n\x17VaccuumQueuedJobsResult"\x14\n\x12\x43heckHealthRequest"$\n\x11\x43heckHealthResult\x12\x0f\n\x07success\x18\x01 \x01(\x08"6\n\x15ListQueuedJobsRequest\x12\x0e\n\x06offset\x18\x01 \x01(\r\x12\r\n\x05limit\x18\x02 \x01(\r"\xf4\x01\n\x14ListQueuedJobsResult\x12\x43\n\x0bqueued_jobs\x18\x01 \x03(\x0b\x32..ListQueuedJobsResult.ListQueuedJobsResultItem\x1a\x96\x01\n\x18ListQueuedJobsResultItem\x12\x0c\n\x04uuid\x18\x01 \x01(\t\x12\x14\n\x0cjob_ref_name\x18\x02 \x01(\t\x12\x14\n\x0cnum_attempts\x18\x03 \x01(\x03\x12\x30\n\x0csubmitted_at\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0e\n\x06status\x18\x05 \x01(\t2\xc0\x02\n\rCommandServer\x12\x35\n\nEnqueueJob\x12\x12.EnqueueJobRequest\x1a\x11.EnqueueJobResult"\x00\x12/\n\x08RetryJob\x12\x10.RetryJobRequest\x1a\x0f.RetryJobResult"\x00\x12\x38\n\x0b\x43heckHealth\x12\x13.CheckHealthRequest\x1a\x12.CheckHealthResult"\x00\x12\x41\n\x0eListQueuedJobs\x12\x16.ListQueuedJobsRequest\x1a\x15.ListQueuedJobsResult"\x00\x12J\n\x11VaccuumQueuedJobs\x12\x19.VaccuumQueuedJobsRequest\x1a\x18.VaccuumQueuedJobsResult"\x00\x62\x06proto3'
22
22
  )
23
23
 
24
24
  _globals = globals()
@@ -34,16 +34,24 @@ if _descriptor._USE_C_DESCRIPTORS == False:
34
34
  _globals["_ENQUEUEJOBREQUEST"]._serialized_end = 187
35
35
  _globals["_ENQUEUEJOBRESULT"]._serialized_start = 189
36
36
  _globals["_ENQUEUEJOBRESULT"]._serialized_end = 261
37
- _globals["_CHECKHEALTHREQUEST"]._serialized_start = 263
38
- _globals["_CHECKHEALTHREQUEST"]._serialized_end = 283
39
- _globals["_CHECKHEALTHRESULT"]._serialized_start = 285
40
- _globals["_CHECKHEALTHRESULT"]._serialized_end = 321
41
- _globals["_LISTQUEUEDJOBSREQUEST"]._serialized_start = 323
42
- _globals["_LISTQUEUEDJOBSREQUEST"]._serialized_end = 377
43
- _globals["_LISTQUEUEDJOBSRESULT"]._serialized_start = 380
44
- _globals["_LISTQUEUEDJOBSRESULT"]._serialized_end = 608
45
- _globals["_LISTQUEUEDJOBSRESULT_LISTQUEUEDJOBSRESULTITEM"]._serialized_start = 474
46
- _globals["_LISTQUEUEDJOBSRESULT_LISTQUEUEDJOBSRESULTITEM"]._serialized_end = 608
47
- _globals["_COMMANDSERVER"]._serialized_start = 611
48
- _globals["_COMMANDSERVER"]._serialized_end = 806
37
+ _globals["_RETRYJOBREQUEST"]._serialized_start = 263
38
+ _globals["_RETRYJOBREQUEST"]._serialized_end = 294
39
+ _globals["_RETRYJOBRESULT"]._serialized_start = 296
40
+ _globals["_RETRYJOBRESULT"]._serialized_end = 366
41
+ _globals["_VACCUUMQUEUEDJOBSREQUEST"]._serialized_start = 368
42
+ _globals["_VACCUUMQUEUEDJOBSREQUEST"]._serialized_end = 394
43
+ _globals["_VACCUUMQUEUEDJOBSRESULT"]._serialized_start = 396
44
+ _globals["_VACCUUMQUEUEDJOBSRESULT"]._serialized_end = 421
45
+ _globals["_CHECKHEALTHREQUEST"]._serialized_start = 423
46
+ _globals["_CHECKHEALTHREQUEST"]._serialized_end = 443
47
+ _globals["_CHECKHEALTHRESULT"]._serialized_start = 445
48
+ _globals["_CHECKHEALTHRESULT"]._serialized_end = 481
49
+ _globals["_LISTQUEUEDJOBSREQUEST"]._serialized_start = 483
50
+ _globals["_LISTQUEUEDJOBSREQUEST"]._serialized_end = 537
51
+ _globals["_LISTQUEUEDJOBSRESULT"]._serialized_start = 540
52
+ _globals["_LISTQUEUEDJOBSRESULT"]._serialized_end = 784
53
+ _globals["_LISTQUEUEDJOBSRESULT_LISTQUEUEDJOBSRESULTITEM"]._serialized_start = 634
54
+ _globals["_LISTQUEUEDJOBSRESULT_LISTQUEUEDJOBSRESULTITEM"]._serialized_end = 784
55
+ _globals["_COMMANDSERVER"]._serialized_start = 787
56
+ _globals["_COMMANDSERVER"]._serialized_end = 1107
49
57
  # @@protoc_insertion_point(module_scope)
@@ -35,6 +35,30 @@ class EnqueueJobResult(_message.Message):
35
35
  self, successfully_queued: bool = ..., queued_job_uuid: _Optional[str] = ...
36
36
  ) -> None: ...
37
37
 
38
+ class RetryJobRequest(_message.Message):
39
+ __slots__ = ("uuid",)
40
+ UUID_FIELD_NUMBER: _ClassVar[int]
41
+ uuid: str
42
+ def __init__(self, uuid: _Optional[str] = ...) -> None: ...
43
+
44
+ class RetryJobResult(_message.Message):
45
+ __slots__ = ("successfully_queued", "queued_job_uuid")
46
+ SUCCESSFULLY_QUEUED_FIELD_NUMBER: _ClassVar[int]
47
+ QUEUED_JOB_UUID_FIELD_NUMBER: _ClassVar[int]
48
+ successfully_queued: bool
49
+ queued_job_uuid: str
50
+ def __init__(
51
+ self, successfully_queued: bool = ..., queued_job_uuid: _Optional[str] = ...
52
+ ) -> None: ...
53
+
54
+ class VaccuumQueuedJobsRequest(_message.Message):
55
+ __slots__ = ()
56
+ def __init__(self) -> None: ...
57
+
58
+ class VaccuumQueuedJobsResult(_message.Message):
59
+ __slots__ = ()
60
+ def __init__(self) -> None: ...
61
+
38
62
  class CheckHealthRequest(_message.Message):
39
63
  __slots__ = ()
40
64
  def __init__(self) -> None: ...
@@ -58,21 +82,24 @@ class ListQueuedJobsRequest(_message.Message):
58
82
  class ListQueuedJobsResult(_message.Message):
59
83
  __slots__ = ("queued_jobs",)
60
84
  class ListQueuedJobsResultItem(_message.Message):
61
- __slots__ = ("uuid", "job_ref_name", "num_attempts", "submitted_at")
85
+ __slots__ = ("uuid", "job_ref_name", "num_attempts", "submitted_at", "status")
62
86
  UUID_FIELD_NUMBER: _ClassVar[int]
63
87
  JOB_REF_NAME_FIELD_NUMBER: _ClassVar[int]
64
88
  NUM_ATTEMPTS_FIELD_NUMBER: _ClassVar[int]
65
89
  SUBMITTED_AT_FIELD_NUMBER: _ClassVar[int]
90
+ STATUS_FIELD_NUMBER: _ClassVar[int]
66
91
  uuid: str
67
92
  job_ref_name: str
68
93
  num_attempts: int
69
94
  submitted_at: _timestamp_pb2.Timestamp
95
+ status: str
70
96
  def __init__(
71
97
  self,
72
98
  uuid: _Optional[str] = ...,
73
99
  job_ref_name: _Optional[str] = ...,
74
100
  num_attempts: _Optional[int] = ...,
75
101
  submitted_at: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ...,
102
+ status: _Optional[str] = ...,
76
103
  ) -> None: ...
77
104
 
78
105
  QUEUED_JOBS_FIELD_NUMBER: _ClassVar[int]
@@ -24,6 +24,11 @@ class CommandServerStub(object):
24
24
  request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobRequest.SerializeToString,
25
25
  response_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobResult.FromString,
26
26
  )
27
+ self.RetryJob = channel.unary_unary(
28
+ "/CommandServer/RetryJob",
29
+ request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.RetryJobRequest.SerializeToString,
30
+ response_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.RetryJobResult.FromString,
31
+ )
27
32
  self.CheckHealth = channel.unary_unary(
28
33
  "/CommandServer/CheckHealth",
29
34
  request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CheckHealthRequest.SerializeToString,
@@ -34,6 +39,11 @@ class CommandServerStub(object):
34
39
  request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.ListQueuedJobsRequest.SerializeToString,
35
40
  response_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.ListQueuedJobsResult.FromString,
36
41
  )
42
+ self.VaccuumQueuedJobs = channel.unary_unary(
43
+ "/CommandServer/VaccuumQueuedJobs",
44
+ request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsRequest.SerializeToString,
45
+ response_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsResult.FromString,
46
+ )
37
47
 
38
48
 
39
49
  class CommandServerServicer(object):
@@ -45,6 +55,12 @@ class CommandServerServicer(object):
45
55
  context.set_details("Method not implemented!")
46
56
  raise NotImplementedError("Method not implemented!")
47
57
 
58
+ def RetryJob(self, request, context):
59
+ """Missing associated documentation comment in .proto file."""
60
+ context.set_code(grpc.StatusCode.UNIMPLEMENTED)
61
+ context.set_details("Method not implemented!")
62
+ raise NotImplementedError("Method not implemented!")
63
+
48
64
  def CheckHealth(self, request, context):
49
65
  """Missing associated documentation comment in .proto file."""
50
66
  context.set_code(grpc.StatusCode.UNIMPLEMENTED)
@@ -57,6 +73,12 @@ class CommandServerServicer(object):
57
73
  context.set_details("Method not implemented!")
58
74
  raise NotImplementedError("Method not implemented!")
59
75
 
76
+ def VaccuumQueuedJobs(self, request, context):
77
+ """Missing associated documentation comment in .proto file."""
78
+ context.set_code(grpc.StatusCode.UNIMPLEMENTED)
79
+ context.set_details("Method not implemented!")
80
+ raise NotImplementedError("Method not implemented!")
81
+
60
82
 
61
83
  def add_CommandServerServicer_to_server(servicer, server):
62
84
  rpc_method_handlers = {
@@ -65,6 +87,11 @@ def add_CommandServerServicer_to_server(servicer, server):
65
87
  request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobRequest.FromString,
66
88
  response_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobResult.SerializeToString,
67
89
  ),
90
+ "RetryJob": grpc.unary_unary_rpc_method_handler(
91
+ servicer.RetryJob,
92
+ request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.RetryJobRequest.FromString,
93
+ response_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.RetryJobResult.SerializeToString,
94
+ ),
68
95
  "CheckHealth": grpc.unary_unary_rpc_method_handler(
69
96
  servicer.CheckHealth,
70
97
  request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CheckHealthRequest.FromString,
@@ -75,6 +102,11 @@ def add_CommandServerServicer_to_server(servicer, server):
75
102
  request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.ListQueuedJobsRequest.FromString,
76
103
  response_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.ListQueuedJobsResult.SerializeToString,
77
104
  ),
105
+ "VaccuumQueuedJobs": grpc.unary_unary_rpc_method_handler(
106
+ servicer.VaccuumQueuedJobs,
107
+ request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsRequest.FromString,
108
+ response_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsResult.SerializeToString,
109
+ ),
78
110
  }
79
111
  generic_handler = grpc.method_handlers_generic_handler(
80
112
  "CommandServer", rpc_method_handlers
@@ -115,6 +147,35 @@ class CommandServer(object):
115
147
  metadata,
116
148
  )
117
149
 
150
+ @staticmethod
151
+ def RetryJob(
152
+ request,
153
+ target,
154
+ options=(),
155
+ channel_credentials=None,
156
+ call_credentials=None,
157
+ insecure=False,
158
+ compression=None,
159
+ wait_for_ready=None,
160
+ timeout=None,
161
+ metadata=None,
162
+ ):
163
+ return grpc.experimental.unary_unary(
164
+ request,
165
+ target,
166
+ "/CommandServer/RetryJob",
167
+ uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.RetryJobRequest.SerializeToString,
168
+ uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.RetryJobResult.FromString,
169
+ options,
170
+ channel_credentials,
171
+ insecure,
172
+ call_credentials,
173
+ compression,
174
+ wait_for_ready,
175
+ timeout,
176
+ metadata,
177
+ )
178
+
118
179
  @staticmethod
119
180
  def CheckHealth(
120
181
  request,
@@ -172,3 +233,32 @@ class CommandServer(object):
172
233
  timeout,
173
234
  metadata,
174
235
  )
236
+
237
+ @staticmethod
238
+ def VaccuumQueuedJobs(
239
+ request,
240
+ target,
241
+ options=(),
242
+ channel_credentials=None,
243
+ call_credentials=None,
244
+ insecure=False,
245
+ compression=None,
246
+ wait_for_ready=None,
247
+ timeout=None,
248
+ metadata=None,
249
+ ):
250
+ return grpc.experimental.unary_unary(
251
+ request,
252
+ target,
253
+ "/CommandServer/VaccuumQueuedJobs",
254
+ uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsRequest.SerializeToString,
255
+ uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsResult.FromString,
256
+ options,
257
+ channel_credentials,
258
+ insecure,
259
+ call_credentials,
260
+ compression,
261
+ wait_for_ready,
262
+ timeout,
263
+ metadata,
264
+ )
@@ -8,6 +8,8 @@ from uncountable.types import queued_job_t
8
8
 
9
9
  class CommandType(StrEnum):
10
10
  ENQUEUE_JOB = "enqueue_job"
11
+ RETRY_JOB = "retry_job"
12
+ VACCUUM_QUEUED_JOBS = "vaccuum_queued_jobs"
11
13
 
12
14
 
13
15
  RT = typing.TypeVar("RT")
@@ -24,6 +26,16 @@ class CommandEnqueueJobResponse:
24
26
  queued_job_uuid: str
25
27
 
26
28
 
29
+ @dataclass(kw_only=True)
30
+ class CommandRetryJobResponse:
31
+ queued_job_uuid: str | None
32
+
33
+
34
+ @dataclass(kw_only=True)
35
+ class CommandVaccuumQueuedJobsResponse:
36
+ pass
37
+
38
+
27
39
  @dataclass(kw_only=True)
28
40
  class CommandEnqueueJob(CommandBase[CommandEnqueueJobResponse]):
29
41
  type: CommandType = CommandType.ENQUEUE_JOB
@@ -32,7 +44,18 @@ class CommandEnqueueJob(CommandBase[CommandEnqueueJobResponse]):
32
44
  response_queue: asyncio.Queue[CommandEnqueueJobResponse]
33
45
 
34
46
 
35
- _Command = CommandEnqueueJob
47
+ @dataclass(kw_only=True)
48
+ class CommandRetryJob(CommandBase[CommandRetryJobResponse]):
49
+ type: CommandType = CommandType.RETRY_JOB
50
+ queued_job_uuid: str
51
+
52
+
53
+ @dataclass(kw_only=True)
54
+ class CommandVaccuumQueuedJobs(CommandBase[CommandVaccuumQueuedJobsResponse]):
55
+ type: CommandType = CommandType.VACCUUM_QUEUED_JOBS
56
+
57
+
58
+ _Command = CommandEnqueueJob | CommandRetryJob | CommandVaccuumQueuedJobs
36
59
 
37
60
 
38
61
  CommandQueue = asyncio.Queue[_Command]
@@ -2,7 +2,7 @@ import datetime
2
2
  import uuid
3
3
  from datetime import UTC
4
4
 
5
- from sqlalchemy import delete, insert, select, update
5
+ from sqlalchemy import delete, insert, or_, select, text, update
6
6
  from sqlalchemy.engine import Engine
7
7
 
8
8
  from pkgs.argument_parser import CachedParser
@@ -14,6 +14,8 @@ from uncountable.types import queued_job_t
14
14
 
15
15
  queued_job_payload_parser = CachedParser(queued_job_t.QueuedJobPayload)
16
16
 
17
+ MAX_QUEUE_WINDOW_DAYS = 30
18
+
17
19
 
18
20
  class DatastoreSqlite(Datastore):
19
21
  def __init__(self, session_maker: DBSessionMaker) -> None:
@@ -23,6 +25,17 @@ class DatastoreSqlite(Datastore):
23
25
  @classmethod
24
26
  def setup(cls, engine: Engine) -> None:
25
27
  Base.metadata.create_all(engine)
28
+ with engine.connect() as connection:
29
+ if not bool(
30
+ connection.execute(
31
+ text(
32
+ "select exists (select 1 from pragma_table_info('queued_jobs') where name='status');"
33
+ )
34
+ ).scalar()
35
+ ):
36
+ connection.execute(
37
+ text("alter table queued_jobs add column status VARCHAR")
38
+ )
26
39
 
27
40
  def add_job_to_queue(
28
41
  self, job_payload: queued_job_t.QueuedJobPayload, job_ref_name: str
@@ -36,6 +49,7 @@ class DatastoreSqlite(Datastore):
36
49
  QueuedJob.id.key: queued_job_uuid,
37
50
  QueuedJob.job_ref_name.key: job_ref_name,
38
51
  QueuedJob.payload.key: serialized_payload,
52
+ QueuedJob.status.key: queued_job_t.JobStatus.QUEUED,
39
53
  QueuedJob.num_attempts: num_attempts,
40
54
  QueuedJob.submitted_at: submitted_at,
41
55
  })
@@ -44,10 +58,48 @@ class DatastoreSqlite(Datastore):
44
58
  queued_job_uuid=queued_job_uuid,
45
59
  job_ref_name=job_ref_name,
46
60
  payload=job_payload,
61
+ status=queued_job_t.JobStatus.QUEUED,
47
62
  submitted_at=submitted_at,
48
63
  num_attempts=num_attempts,
49
64
  )
50
65
 
66
+ def retry_job(
67
+ self,
68
+ queued_job_uuid: str,
69
+ ) -> queued_job_t.QueuedJob | None:
70
+ with self.session_maker() as session:
71
+ select_stmt = select(
72
+ QueuedJob.id,
73
+ QueuedJob.payload,
74
+ QueuedJob.num_attempts,
75
+ QueuedJob.job_ref_name,
76
+ QueuedJob.status,
77
+ QueuedJob.submitted_at,
78
+ ).filter(QueuedJob.id == queued_job_uuid)
79
+ existing_job = session.execute(select_stmt).one_or_none()
80
+
81
+ if (
82
+ existing_job is None
83
+ or existing_job.status != queued_job_t.JobStatus.FAILED
84
+ ):
85
+ return None
86
+
87
+ update_stmt = (
88
+ update(QueuedJob)
89
+ .values({QueuedJob.status.key: queued_job_t.JobStatus.QUEUED})
90
+ .filter(QueuedJob.id == queued_job_uuid)
91
+ )
92
+ session.execute(update_stmt)
93
+
94
+ return queued_job_t.QueuedJob(
95
+ queued_job_uuid=existing_job.id,
96
+ job_ref_name=existing_job.job_ref_name,
97
+ num_attempts=existing_job.num_attempts,
98
+ status=queued_job_t.JobStatus.QUEUED,
99
+ submitted_at=existing_job.submitted_at,
100
+ payload=queued_job_payload_parser.parse_storage(existing_job.payload),
101
+ )
102
+
51
103
  def increment_num_attempts(self, queued_job_uuid: str) -> int:
52
104
  with self.session_maker() as session:
53
105
  update_stmt = (
@@ -68,6 +120,17 @@ class DatastoreSqlite(Datastore):
68
120
  delete_stmt = delete(QueuedJob).filter(QueuedJob.id == queued_job_uuid)
69
121
  session.execute(delete_stmt)
70
122
 
123
+ def update_job_status(
124
+ self, queued_job_uuid: str, status: queued_job_t.JobStatus
125
+ ) -> None:
126
+ with self.session_maker() as session:
127
+ update_stmt = (
128
+ update(QueuedJob)
129
+ .values({QueuedJob.status.key: status})
130
+ .filter(QueuedJob.id == queued_job_uuid)
131
+ )
132
+ session.execute(update_stmt)
133
+
71
134
  def list_queued_job_metadata(
72
135
  self, offset: int = 0, limit: int | None = 100
73
136
  ) -> list[queued_job_t.QueuedJobMetadata]:
@@ -77,6 +140,7 @@ class DatastoreSqlite(Datastore):
77
140
  QueuedJob.id,
78
141
  QueuedJob.job_ref_name,
79
142
  QueuedJob.num_attempts,
143
+ QueuedJob.status,
80
144
  QueuedJob.submitted_at,
81
145
  )
82
146
  .order_by(QueuedJob.submitted_at)
@@ -89,6 +153,7 @@ class DatastoreSqlite(Datastore):
89
153
  queued_job_uuid=row.id,
90
154
  job_ref_name=row.job_ref_name,
91
155
  num_attempts=row.num_attempts,
156
+ status=row.status or queued_job_t.JobStatus.QUEUED,
92
157
  submitted_at=row.submitted_at,
93
158
  )
94
159
  for row in session.execute(select_statement)
@@ -106,9 +171,16 @@ class DatastoreSqlite(Datastore):
106
171
  QueuedJob.payload,
107
172
  QueuedJob.num_attempts,
108
173
  QueuedJob.job_ref_name,
174
+ QueuedJob.status,
109
175
  QueuedJob.submitted_at,
110
176
  )
111
177
  .filter(QueuedJob.job_ref_name == job_ref_name)
178
+ .filter(
179
+ or_(
180
+ QueuedJob.status == queued_job_t.JobStatus.QUEUED,
181
+ QueuedJob.status.is_(None),
182
+ )
183
+ )
112
184
  .limit(1)
113
185
  .order_by(QueuedJob.submitted_at)
114
186
  )
@@ -119,6 +191,7 @@ class DatastoreSqlite(Datastore):
119
191
  queued_job_uuid=row.id,
120
192
  job_ref_name=row.job_ref_name,
121
193
  num_attempts=row.num_attempts,
194
+ status=row.status or queued_job_t.JobStatus.QUEUED,
122
195
  submitted_at=row.submitted_at,
123
196
  payload=parsed_payload,
124
197
  )
@@ -127,13 +200,23 @@ class DatastoreSqlite(Datastore):
127
200
 
128
201
  def load_job_queue(self) -> list[queued_job_t.QueuedJob]:
129
202
  with self.session_maker() as session:
130
- select_stmt = select(
131
- QueuedJob.id,
132
- QueuedJob.payload,
133
- QueuedJob.num_attempts,
134
- QueuedJob.job_ref_name,
135
- QueuedJob.submitted_at,
136
- ).order_by(QueuedJob.submitted_at)
203
+ select_stmt = (
204
+ select(
205
+ QueuedJob.id,
206
+ QueuedJob.payload,
207
+ QueuedJob.num_attempts,
208
+ QueuedJob.job_ref_name,
209
+ QueuedJob.status,
210
+ QueuedJob.submitted_at,
211
+ )
212
+ .filter(
213
+ or_(
214
+ QueuedJob.status == queued_job_t.JobStatus.QUEUED,
215
+ QueuedJob.status.is_(None),
216
+ )
217
+ )
218
+ .order_by(QueuedJob.submitted_at)
219
+ )
137
220
 
138
221
  queued_jobs: list[queued_job_t.QueuedJob] = []
139
222
  for row in session.execute(select_stmt):
@@ -143,9 +226,25 @@ class DatastoreSqlite(Datastore):
143
226
  queued_job_uuid=row.id,
144
227
  job_ref_name=row.job_ref_name,
145
228
  num_attempts=row.num_attempts,
229
+ status=row.status or queued_job_t.JobStatus.QUEUED,
146
230
  submitted_at=row.submitted_at,
147
231
  payload=parsed_payload,
148
232
  )
149
233
  )
150
234
 
151
235
  return queued_jobs
236
+
237
+ def vaccuum_queued_jobs(self) -> None:
238
+ with self.session_maker() as session:
239
+ delete_stmt = (
240
+ delete(QueuedJob)
241
+ .filter(QueuedJob.status == queued_job_t.JobStatus.QUEUED)
242
+ .filter(
243
+ QueuedJob.submitted_at
244
+ <= (
245
+ datetime.datetime.now(UTC)
246
+ - datetime.timedelta(days=MAX_QUEUE_WINDOW_DAYS)
247
+ )
248
+ )
249
+ )
250
+ session.execute(delete_stmt)
@@ -1,7 +1,9 @@
1
- from sqlalchemy import JSON, BigInteger, Column, DateTime, Text
1
+ from sqlalchemy import JSON, BigInteger, Column, DateTime, Enum, Text
2
2
  from sqlalchemy.orm import declarative_base
3
3
  from sqlalchemy.sql import func
4
4
 
5
+ from uncountable.types import queued_job_t
6
+
5
7
  Base = declarative_base()
6
8
 
7
9
 
@@ -15,3 +17,8 @@ class QueuedJob(Base):
15
17
  )
16
18
  payload = Column(JSON, nullable=False)
17
19
  num_attempts = Column(BigInteger, nullable=False, default=0, server_default="0")
20
+ status = Column(
21
+ Enum(queued_job_t.JobStatus, length=None),
22
+ default=queued_job_t.JobStatus.QUEUED,
23
+ nullable=True,
24
+ )