UncountablePythonSDK 0.0.133__py3-none-any.whl → 0.0.136.dev0__py3-none-any.whl

uncountable/integration/queue_runner/worker.py

@@ -87,30 +87,32 @@ def run_queued_job(
             base_span=span,
             profile_metadata=job_details.profile_metadata,
             job_definition=job_details.job_definition,
+            queued_job_uuid=queued_job.queued_job_uuid,
         )
-        try:
-            client = construct_uncountable_client(
-                profile_meta=job_details.profile_metadata, logger=job_logger
-            )
-            batch_processor = AsyncBatchProcessor(client=client)
-
-            payload = _resolve_queued_job_payload(queued_job)
-
-            args = JobArguments(
-                job_definition=job_details.job_definition,
-                client=client,
-                batch_processor=batch_processor,
-                profile_metadata=job_details.profile_metadata,
-                logger=job_logger,
-                payload=payload,
-                job_uuid=queued_job.queued_job_uuid,
-            )
-
-            return execute_job(
-                args=args,
-                profile_metadata=job_details.profile_metadata,
-                job_definition=job_details.job_definition,
-            )
-        except BaseException as e:
-            job_logger.log_exception(e)
-            return job_definition_t.JobResult(success=False)
+        with job_logger.resource_tracking():
+            try:
+                client = construct_uncountable_client(
+                    profile_meta=job_details.profile_metadata, logger=job_logger
+                )
+                batch_processor = AsyncBatchProcessor(client=client)
+
+                payload = _resolve_queued_job_payload(queued_job)
+
+                args = JobArguments(
+                    job_definition=job_details.job_definition,
+                    client=client,
+                    batch_processor=batch_processor,
+                    profile_metadata=job_details.profile_metadata,
+                    logger=job_logger,
+                    payload=payload,
+                    job_uuid=queued_job.queued_job_uuid,
+                )
+
+                return execute_job(
+                    args=args,
+                    profile_metadata=job_details.profile_metadata,
+                    job_definition=job_details.job_definition,
+                )
+            except BaseException as e:
+                job_logger.log_exception(e)
+                return job_definition_t.JobResult(success=False)

uncountable/integration/scheduler.py

@@ -6,6 +6,7 @@ import time
 from dataclasses import dataclass
 from datetime import UTC
 from enum import StrEnum
+from typing import assert_never
 
 from opentelemetry.trace import get_current_span
 
@@ -16,6 +17,7 @@ from uncountable.integration.queue_runner.command_server import (
     check_health,
 )
 from uncountable.integration.queue_runner.queue_runner import start_queue_runner
+from uncountable.integration.queue_runner.types import RESTART_EXIT_CODE
 from uncountable.integration.telemetry import Logger
 
 SHUTDOWN_TIMEOUT_SECS = 30
@@ -55,6 +57,19 @@ class ProcessInfo:
                 return self.process.poll()
 
 
+@dataclass(kw_only=True)
+class ProcessAlarmRestart:
+    process: ProcessInfo
+
+
+@dataclass(kw_only=True)
+class ProcessAlarmShutdownAll:
+    pass
+
+
+ProcessAlarm = ProcessAlarmRestart | ProcessAlarmShutdownAll
+
+
 def handle_shutdown(logger: Logger, processes: dict[ProcessName, ProcessInfo]) -> None:
     logger.log_info("received shutdown command, shutting down sub-processes")
     for proc_info in processes.values():
@@ -129,12 +144,21 @@ def restart_process(
             logger.log_info("uwsgi restarted successfully")
 
 
-def check_process_alive(
+def check_process_alarms(
     logger: Logger, processes: dict[ProcessName, ProcessInfo]
-) -> None:
+) -> ProcessAlarm | None:
     for proc_info in processes.values():
         if not proc_info.is_alive:
-            restart_process(logger, proc_info, processes)
+            if proc_info.exitcode == RESTART_EXIT_CODE:
+                logger.log_warning(
+                    f"process {proc_info.name} requested restart! restarting"
+                )
+                return ProcessAlarmRestart(process=proc_info)
+            logger.log_error(
+                f"process {proc_info.name} shut down unexpectedly! shutting down scheduler; exit code is {proc_info.exitcode}"
+            )
+            return ProcessAlarmShutdownAll()
+    return None
 
 
 def _wait_queue_runner_online() -> None:
@@ -166,16 +190,24 @@ def main() -> None:
         processes[process.name] = process
         logger.log_info(f"started process {process.name}")
 
-    runner_process = multiprocessing.Process(target=start_queue_runner)
-    runner_process.start()
-    add_process(ProcessInfo(name=ProcessName.QUEUE_RUNNER, process=runner_process))
+    def _start_queue_runner() -> None:
+        runner_process = multiprocessing.Process(target=start_queue_runner)
+        runner_process.start()
+        add_process(
+            ProcessInfo(
+                name=ProcessName.QUEUE_RUNNER,
+                process=runner_process,
+            )
+        )
 
-    try:
-        _wait_queue_runner_online()
-    except Exception as e:
-        logger.log_exception(e)
-        handle_shutdown(logger, processes=processes)
-        return
+        try:
+            _wait_queue_runner_online()
+        except Exception as e:
+            logger.log_exception(e)
+            handle_shutdown(logger, processes=processes)
+            return
+
+    _start_queue_runner()
 
     cron_process = multiprocessing.Process(target=cron_target)
     cron_process.start()
@@ -189,7 +221,26 @@ def main() -> None:
 
     try:
         while True:
-            check_process_alive(logger, processes=processes)
+            process_alarm = check_process_alarms(logger, processes=processes)
+            match process_alarm:
+                case ProcessAlarmRestart():
+                    match process_alarm.process.name:
+                        case ProcessName.QUEUE_RUNNER:
+                            del processes[ProcessName.QUEUE_RUNNER]
+                            _start_queue_runner()
+                        case ProcessName.CRON_SERVER | ProcessName.UWSGI:
+                            raise NotImplementedError(
+                                f"restarting {process_alarm.process.name} not yet implemented"
+                            )
+                        case _:
+                            assert_never(process_alarm.process.name)
+                case ProcessAlarmShutdownAll():
+                    handle_shutdown(logger, processes)
+                    sys.exit(1)
+                case None:
+                    pass
+                case _:
+                    assert_never(process_alarm)
             time.sleep(1)
     except KeyboardInterrupt:
         handle_shutdown(logger, processes=processes)

uncountable/integration/telemetry.py

@@ -1,13 +1,16 @@
 import functools
 import json
 import os
+import threading
 import time
 import traceback
+import types
 import typing
 from contextlib import contextmanager
 from enum import StrEnum
 from typing import Generator, assert_never, cast
 
+import psutil
 from opentelemetry import _logs, trace
 from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
@@ -185,6 +188,76 @@ class Logger:
             yield self
 
 
+class PerJobResourceTracker:
+    def __init__(self, logger: "JobLogger", sample_interval: float = 0.5) -> None:
+        self.logger = logger
+        self.sample_interval = sample_interval
+        self._process = psutil.Process(os.getpid())
+        self._stop_event = threading.Event()
+        self._thread: threading.Thread | None = None
+
+        self.max_rss: int = 0
+        self.start_cpu_times: psutil._common.pcputimes | None = None
+        self.end_cpu_times: psutil._common.pcputimes | None = None
+        self.start_wall_time: float | None = None
+        self.end_wall_time: float | None = None
+
+    def start(self) -> None:
+        self.start_cpu_times = self._process.cpu_times()
+        self.start_wall_time = time.monotonic()
+
+        def _monitor() -> None:
+            try:
+                while not self._stop_event.is_set():
+                    rss = self._process.memory_info().rss
+                    self.max_rss = max(self.max_rss, rss)
+                    time.sleep(self.sample_interval)
+            except Exception:
+                self._stop_event.set()
+
+        self._thread = threading.Thread(target=_monitor, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop_event.set()
+        if self._thread is not None:
+            self._thread.join()
+        self.end_cpu_times = self._process.cpu_times()
+        self.end_wall_time = time.monotonic()
+
+    def __enter__(self) -> typing.Self:
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> None:
+        self.stop()
+        stats = dict(self.summary())
+        self.logger.log_info("Job resource usage summary", attributes=stats)
+
+    def summary(self) -> Attributes:
+        assert self.start_cpu_times is not None
+        assert self.end_cpu_times is not None
+        assert self.start_wall_time is not None
+        assert self.end_wall_time is not None
+
+        cpu_user = self.end_cpu_times.user - self.start_cpu_times.user
+        cpu_sys = self.end_cpu_times.system - self.start_cpu_times.system
+        cpu_total = cpu_user + cpu_sys
+        elapsed = self.end_wall_time - self.start_wall_time
+        return {
+            "cpu_user_s": round(cpu_user, 3),
+            "cpu_system_s": round(cpu_sys, 3),
+            "cpu_total_s": round(cpu_total, 3),
+            "wall_time_s": round(elapsed, 3),
+            "peak_rss_mb": round(self.max_rss / (1024 * 1024), 2),
+        }
+
+
 class JobLogger(Logger):
     def __init__(
         self,
@@ -192,9 +265,11 @@ class JobLogger(Logger):
         base_span: Span,
         profile_metadata: job_definition_t.ProfileMetadata,
         job_definition: job_definition_t.JobDefinition,
+        queued_job_uuid: str,
     ) -> None:
         self.profile_metadata = profile_metadata
         self.job_definition = job_definition
+        self.queued_job_uuid = queued_job_uuid
         super().__init__(base_span)
 
     def _patch_attributes(
@@ -214,6 +289,7 @@ class JobLogger(Logger):
         patched_attributes["job.name"] = self.job_definition.name
         patched_attributes["job.id"] = self.job_definition.id
         patched_attributes["job.definition_type"] = self.job_definition.type
+        patched_attributes["job.queued_job_uuid"] = self.queued_job_uuid
         match self.job_definition:
             case job_definition_t.CronJobDefinition():
                 patched_attributes["job.definition.cron_spec"] = (
@@ -239,6 +315,9 @@ class JobLogger(Logger):
                 assert_never(self.job_definition.executor)
         return _cast_attributes(patched_attributes)
 
+    def resource_tracking(self) -> PerJobResourceTracker:
+        return PerJobResourceTracker(self)
+
 
 @contextmanager
 def push_scope_optional(

uncountable/types/__init__.py

@@ -37,6 +37,7 @@ from . import entity_t as entity_t
 from .api.batch import execute_batch as execute_batch_t
 from .api.batch import execute_batch_load_async as execute_batch_load_async_t
 from . import experiment_groups_t as experiment_groups_t
+from .api.runsheet import export_default_runsheet as export_default_runsheet_t
 from .api.entity import export_entities as export_entities_t
 from . import exports_t as exports_t
 from .api.listing import fetch_listing as fetch_listing_t
@@ -121,7 +122,6 @@ from .api.recipes import set_recipe_tags as set_recipe_tags_t
 from .api.recipes import set_recipe_total as set_recipe_total_t
 from .api.entity import set_values as set_values_t
 from . import sockets_t as sockets_t
-from . import structured_filters_t as structured_filters_t
 from .api.entity import transition_entity_phase as transition_entity_phase_t
 from .api.recipes import unarchive_recipes as unarchive_recipes_t
 from . import units_t as units_t
@@ -172,6 +172,7 @@ __all__: list[str] = [
     "execute_batch_t",
     "execute_batch_load_async_t",
     "experiment_groups_t",
+    "export_default_runsheet_t",
     "export_entities_t",
     "exports_t",
     "fetch_listing_t",
@@ -256,7 +257,6 @@ __all__: list[str] = [
     "set_recipe_total_t",
     "set_values_t",
     "sockets_t",
-    "structured_filters_t",
     "transition_entity_phase_t",
     "unarchive_recipes_t",
     "units_t",

uncountable/types/api/files/download_file.py

@@ -21,6 +21,7 @@ __all__: list[str] = [
     "FileDownloadQuery",
     "FileDownloadQueryBase",
     "FileDownloadQueryEntityField",
+    "FileDownloadQueryTextDocumentId",
     "FileDownloadQueryType",
 ]
 
@@ -31,6 +32,7 @@ ENDPOINT_PATH = "api/external/files/download_file"
 # DO NOT MODIFY -- This file is generated by type_spec
 class FileDownloadQueryType(StrEnum):
     ENTITY_FIELD = "entity_field"
+    TEXT_DOCUMENT_ID = "text_document_id"
 
 
 # DO NOT MODIFY -- This file is generated by type_spec
@@ -54,14 +56,26 @@ class FileDownloadQueryEntityField(FileDownloadQueryBase):
     field_key: identifier_t.IdentifierKey
 
 
+# DO NOT MODIFY -- This file is generated by type_spec
+@serial_class(
+    named_type_path="sdk.api.files.download_file.FileDownloadQueryTextDocumentId",
+    parse_require={"type"},
+)
+@dataclasses.dataclass(slots=base_t.ENABLE_SLOTS, kw_only=True)  # type: ignore[literal-required]
+class FileDownloadQueryTextDocumentId(FileDownloadQueryBase):
+    type: typing.Literal[FileDownloadQueryType.TEXT_DOCUMENT_ID] = FileDownloadQueryType.TEXT_DOCUMENT_ID
+    text_document_id: base_t.ObjectId
+
+
 # DO NOT MODIFY -- This file is generated by type_spec
 FileDownloadQuery = typing.Annotated[
-    typing.Union[FileDownloadQueryEntityField],
+    FileDownloadQueryEntityField | FileDownloadQueryTextDocumentId,
     serial_union_annotation(
     named_type_path="sdk.api.files.download_file.FileDownloadQuery",
         discriminator="type",
         discriminator_map={
             "entity_field": FileDownloadQueryEntityField,
+            "text_document_id": FileDownloadQueryTextDocumentId,
         },
     ),
 ]

uncountable/types/api/listing/fetch_listing.py

@@ -11,7 +11,6 @@ from pkgs.serialization import serial_class
 from ... import base_t
 from ... import entity_t
 from ... import listing_t
-from ... import structured_filters_t
 
 __all__: list[str] = [
     "Arguments",
@@ -33,7 +32,7 @@ ENDPOINT_PATH = "api/external/listing/fetch_listing"
 class Arguments:
     entity_type: entity_t.EntityType
     columns: list[listing_t.ColumnIdentifier]
-    filters: structured_filters_t.FilterNode | None = None
+    filters: listing_t.FilterNode | None = None
     limit: int | None = None
     offset: int | None = None
 

uncountable/types/api/runsheet/export_default_runsheet.py

@@ -0,0 +1,44 @@
+# DO NOT MODIFY -- This file is generated by type_spec
+# ruff: noqa: E402 Q003
+# fmt: off
+# isort: skip_file
+from __future__ import annotations
+import typing  # noqa: F401
+import datetime  # noqa: F401
+from decimal import Decimal  # noqa: F401
+import dataclasses
+from pkgs.serialization import serial_class
+from ... import base_t
+from ... import entity_t
+from ... import identifier_t
+
+__all__: list[str] = [
+    "Arguments",
+    "Data",
+    "ENDPOINT_METHOD",
+    "ENDPOINT_PATH",
+]
+
+ENDPOINT_METHOD = "POST"
+ENDPOINT_PATH = "api/external/runsheet/export_default_runsheet"
+
+
+# DO NOT MODIFY -- This file is generated by type_spec
+@serial_class(
+    named_type_path="sdk.api.runsheet.export_default_runsheet.Arguments",
+)
+@dataclasses.dataclass(slots=base_t.ENABLE_SLOTS, kw_only=True)  # type: ignore[literal-required]
+class Arguments:
+    entities: list[identifier_t.IdentifierKey]
+    entity_type: entity_t.EntityType = entity_t.EntityType.RECIPE
+    runsheet_key: identifier_t.IdentifierKey
+
+
+# DO NOT MODIFY -- This file is generated by type_spec
+@serial_class(
+    named_type_path="sdk.api.runsheet.export_default_runsheet.Data",
+)
+@dataclasses.dataclass(slots=base_t.ENABLE_SLOTS, kw_only=True)  # type: ignore[literal-required]
+class Data:
+    text_document_id: base_t.ObjectId
+# DO NOT MODIFY -- This file is generated by type_spec

uncountable/types/client_base.py

@@ -37,6 +37,7 @@ import uncountable.types.api.batch.execute_batch as execute_batch_t
 import uncountable.types.api.batch.execute_batch_load_async as execute_batch_load_async_t
 import uncountable.types.api.entity.export_entities as export_entities_t
 from uncountable.types import exports_t
+import uncountable.types.api.listing.fetch_listing as fetch_listing_t
 from uncountable.types import field_values_t
 from uncountable.types import generic_upload_t
 import uncountable.types.api.recipes.get_column_calculation_values as get_column_calculation_values_t
@@ -49,6 +50,7 @@ import uncountable.types.api.inputs.get_input_names as get_input_names_t
 import uncountable.types.api.inputs.get_inputs_data as get_inputs_data_t
 import uncountable.types.api.outputs.get_output_data as get_output_data_t
 import uncountable.types.api.outputs.get_output_names as get_output_names_t
+import uncountable.types.api.outputs.get_output_organization as get_output_organization_t
 import uncountable.types.api.project.get_projects as get_projects_t
 import uncountable.types.api.project.get_projects_data as get_projects_data_t
 import uncountable.types.api.recipes.get_recipe_calculations as get_recipe_calculations_t
@@ -64,6 +66,7 @@ import uncountable.types.api.uploader.invoke_uploader as invoke_uploader_t
 import uncountable.types.api.entity.list_aggregate as list_aggregate_t
 import uncountable.types.api.entity.list_entities as list_entities_t
 import uncountable.types.api.id_source.list_id_source as list_id_source_t
+from uncountable.types import listing_t
 import uncountable.types.api.entity.lock_entity as lock_entity_t
 import uncountable.types.api.recipes.lock_recipes as lock_recipes_t
 import uncountable.types.api.entity.lookup_entity as lookup_entity_t
@@ -738,6 +741,39 @@ class ClientMethods(ABC):
         )
         return self.do_request(api_request=api_request, return_type=export_entities_t.Data)
 
+    def fetch_listing(
+        self,
+        *,
+        entity_type: entity_t.EntityType,
+        columns: list[listing_t.ColumnIdentifier],
+        filters: listing_t.FilterNode | None = None,
+        limit: int | None = None,
+        offset: int | None = None,
+        _request_options: client_config_t.RequestOptions | None = None,
+    ) -> fetch_listing_t.Data:
+        """External API for fetching a listing based on some structured columns and filters
+
+        :param entity_type: The entity type to fetch listing entries for
+        :param columns: The columns to include in the results
+        :param filters: Structured filters to apply to the listing, represented by a FilterNode object
+        :param limit: The number of entries to return. If not provided, defaults to 100. Note a maximum of 100 entries can be returned.
+        :param offset: The number of entries to skip for pagination
+        """
+        args = fetch_listing_t.Arguments(
+            entity_type=entity_type,
+            columns=columns,
+            filters=filters,
+            limit=limit,
+            offset=offset,
+        )
+        api_request = APIRequest(
+            method=fetch_listing_t.ENDPOINT_METHOD,
+            endpoint=fetch_listing_t.ENDPOINT_PATH,
+            args=args,
+            request_options=_request_options,
+        )
+        return self.do_request(api_request=api_request, return_type=fetch_listing_t.Data)
+
     def get_column_calculation_values(
         self,
         *,
@@ -989,6 +1025,26 @@ class ClientMethods(ABC):
         )
         return self.do_request(api_request=api_request, return_type=get_output_names_t.Data)
 
+    def get_output_organization(
+        self,
+        *,
+        request: get_output_organization_t.OutputOrganizationRequest,
+        _request_options: client_config_t.RequestOptions | None = None,
+    ) -> get_output_organization_t.Data:
+        """Returns the output organization for the specified material family, project, or user
+
+        """
+        args = get_output_organization_t.Arguments(
+            request=request,
+        )
+        api_request = APIRequest(
+            method=get_output_organization_t.ENDPOINT_METHOD,
+            endpoint=get_output_organization_t.ENDPOINT_PATH,
+            args=args,
+            request_options=_request_options,
+        )
+        return self.do_request(api_request=api_request, return_type=get_output_organization_t.Data)
+
     def get_projects(
         self,
         *,

uncountable/types/listing.py

@@ -5,5 +5,42 @@
 # Kept only for SDK backwards compatibility
 from .listing_t import ColumnType as ColumnType
 from .listing_t import ColumnIdentifierEntityRefName as ColumnIdentifierEntityRefName
+from .listing_t import ColumnIdentifierTransitive as ColumnIdentifierTransitive
+from .listing_t import LoadAggregateType as LoadAggregateType
+from .listing_t import LoadAggregateGroupBy as LoadAggregateGroupBy
+from .listing_t import LoadAggregateCollect as LoadAggregateCollect
+from .listing_t import LoadAggregateCollectDistinct as LoadAggregateCollectDistinct
+from .listing_t import LoadAggregateCount as LoadAggregateCount
+from .listing_t import LoadAggregateCountDistinct as LoadAggregateCountDistinct
+from .listing_t import LoadAggregateFirst as LoadAggregateFirst
+from .listing_t import LoadAggregateUniqueValue as LoadAggregateUniqueValue
+from .listing_t import LoadAggregateMin as LoadAggregateMin
+from .listing_t import LoadAggregateMax as LoadAggregateMax
+from .listing_t import LoadAggregateMean as LoadAggregateMean
+from .listing_t import LoadAggregateStddev as LoadAggregateStddev
+from .listing_t import LoadAggregateSum as LoadAggregateSum
+from .listing_t import DateGroupType as DateGroupType
+from .listing_t import LoadAggregateGroupByDate as LoadAggregateGroupByDate
+from .listing_t import AggregateLoadTypes as AggregateLoadTypes
+from .listing_t import ColumnIdentifierTransitiveAggregate as ColumnIdentifierTransitiveAggregate
 from .listing_t import ColumnIdentifier as ColumnIdentifier
+from .listing_t import FilterRelation as FilterRelation
+from .listing_t import FilterSpecBase as FilterSpecBase
+from .listing_t import FilterScalarType as FilterScalarType
+from .listing_t import FilterIdType as FilterIdType
+from .listing_t import StringFilterValue as StringFilterValue
+from .listing_t import FilterSpecEquals as FilterSpecEquals
+from .listing_t import FilterSpecInclude as FilterSpecInclude
+from .listing_t import FilterSpecIStrContains as FilterSpecIStrContains
+from .listing_t import FilterSpecIStrStartsWith as FilterSpecIStrStartsWith
+from .listing_t import FilterSpecExists as FilterSpecExists
+from .listing_t import FilterSpecGreater as FilterSpecGreater
+from .listing_t import FilterSpecLess as FilterSpecLess
+from .listing_t import FilterSpecGeq as FilterSpecGeq
+from .listing_t import FilterSpecLeq as FilterSpecLeq
+from .listing_t import FilterSpec as FilterSpec
+from .listing_t import FilterNodeType as FilterNodeType
+from .listing_t import FilterNodeBase as FilterNodeBase
+from .listing_t import FilterNodeColumnAnd as FilterNodeColumnAnd
+from .listing_t import FilterNode as FilterNode
 # DO NOT MODIFY -- This file is generated by type_spec