_raylet.pxd 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. # cython: profile=False
  2. # distutils: language = c++
  3. # cython: embedsignature = True
  4. # cython: language_level = 3
  5. from cpython.pystate cimport PyThreadState_Get
  6. from libc.stdint cimport (
  7. int64_t,
  8. )
  9. from libcpp cimport bool as c_bool
  10. from libcpp.string cimport string as c_string
  11. from libcpp.vector cimport vector as c_vector
  12. from libcpp.unordered_map cimport unordered_map
  13. from libcpp.memory cimport (
  14. shared_ptr,
  15. unique_ptr
  16. )
  17. from libcpp.pair cimport pair as c_pair
  18. from libcpp.utility cimport pair
  19. from ray.includes.optional cimport (
  20. optional,
  21. nullopt,
  22. make_optional,
  23. )
  24. from ray.includes.common cimport (
  25. CBuffer,
  26. CRayObject,
  27. CAddress,
  28. CConcurrencyGroup,
  29. CSchedulingStrategy,
  30. CLabelMatchExpressions,
  31. )
  32. from ray.includes.libcoreworker cimport (
  33. ActorHandleSharedPtr,
  34. CActorHandle,
  35. CFiberEvent,
  36. )
  37. from ray.includes.unique_ids cimport (
  38. CObjectID,
  39. CActorID,
  40. CTaskID,
  41. )
  42. from ray.includes.function_descriptor cimport (
  43. CFunctionDescriptor,
  44. )
  45. cdef extern from *:
  46. """
  47. #if __OPTIMIZE__ && __OPTIMIZE__ == 1
  48. #undef __OPTIMIZE__
  49. int __OPTIMIZE__ = 1;
  50. #define __OPTIMIZE__ 1
  51. #elif defined(BAZEL_OPT)
  52. // For compilers that don't define __OPTIMIZE__
  53. int __OPTIMIZE__ = 1;
  54. #else
  55. int __OPTIMIZE__ = 0;
  56. #endif
  57. """
  58. int __OPTIMIZE__
  59. cdef extern from "Python.h":
  60. # Note(simon): This is used to configure asyncio actor stack size.
  61. # Cython made PyThreadState an opaque types. Saying that if the user wants
  62. # specific attributes, they can be declared manually.
  63. # You can find the cpython definition in Include/cpython/pystate.h#L59
  64. ctypedef struct CPyThreadState "PyThreadState":
  65. int recursion_limit
  66. int recursion_remaining
  67. int c_recursion_remaining
  68. # From Include/ceveal.h#67
  69. int Py_GetRecursionLimit()
  70. void Py_SetRecursionLimit(int)
  71. # Note that `functional.pxd` in the Cython repository supports only a limited subset of
  72. # <functional>. Therefore, `from libcpp.functional cimport function` is not enough, and we
  73. # still need to expose some functions here.
  74. cdef extern from "<functional>" namespace "std" nogil:
  75. T bind[T, Args](T callable, Args args)
  76. # Reference: https://github.com/scipy/scipy/blob/6b56162fa6880b0182faea44af88d6a1587f35a8/scipy/stats/_qmc_cy.pyx#L31-L34
  77. cdef cppclass reference_wrapper[T]:
  78. pass
  79. cdef reference_wrapper[T] ref[T](T&)
  80. cdef class Buffer:
  81. cdef:
  82. shared_ptr[CBuffer] buffer
  83. Py_ssize_t shape
  84. Py_ssize_t strides
  85. @staticmethod
  86. cdef make(const shared_ptr[CBuffer]& buffer)
  87. cdef class BaseID:
  88. # To avoid the error of "Python int too large to convert to C ssize_t",
  89. # here `cdef size_t` is required.
  90. cdef size_t hash(self)
  91. cdef class ObjectRef(BaseID):
  92. cdef:
  93. CObjectID data
  94. c_string owner_addr
  95. # Flag indicating whether or not this object ref was added to the set
  96. # of active IDs in the core worker so we know whether we should clean
  97. # it up.
  98. c_bool in_core_worker
  99. c_string call_site_data
  100. # Python object to store optional tensor transport string
  101. object _tensor_transport
  102. cdef CObjectID native(self)
  103. cdef optional[c_string] c_tensor_transport(self)
  104. cdef class ActorID(BaseID):
  105. cdef CActorID data
  106. cdef CActorID native(self)
  107. cdef size_t hash(self)
  108. cdef class CoreWorker:
  109. cdef:
  110. c_bool is_driver
  111. object async_thread
  112. object async_event_loop
  113. object job_config
  114. object current_runtime_env
  115. c_bool is_local_mode
  116. object cgname_to_eventloop_dict
  117. object eventloop_for_default_cg
  118. object thread_for_default_cg
  119. object fd_to_cgname_dict
  120. object _task_id_to_future_lock
  121. dict _task_id_to_future
  122. object event_loop_executor
  123. object _gc_thread
  124. cdef unique_ptr[CAddress] _convert_python_address(self, address=*)
  125. cdef put_serialized_object_and_increment_local_ref(
  126. self,
  127. serialized_object,
  128. optional[c_string] c_tensor_transport,
  129. c_bool pin_object=*,
  130. owner_address=*,
  131. c_bool inline_small_object=*,
  132. c_bool _is_experimental_channel=*,
  133. )
  134. cdef store_task_output(
  135. self, serialized_object,
  136. const CObjectID &return_id,
  137. const CObjectID &generator_id,
  138. size_t data_size, shared_ptr[CBuffer] &metadata, const c_vector[CObjectID]
  139. &contained_id, const CAddress &caller_address,
  140. int64_t *task_output_inlined_bytes,
  141. shared_ptr[CRayObject] *return_ptr)
  142. cdef store_task_outputs(
  143. self,
  144. worker, outputs,
  145. const CAddress &caller_address,
  146. c_vector[c_pair[CObjectID, shared_ptr[CRayObject]]] *returns,
  147. ref_generator_id=*, # CObjectID
  148. optional[c_string] c_tensor_transport=*,
  149. )
  150. cdef make_actor_handle(self, ActorHandleSharedPtr c_actor_handle,
  151. c_bool weak_ref)
  152. cdef c_function_descriptors_to_python(
  153. self, const c_vector[CFunctionDescriptor] &c_function_descriptors)
  154. cdef initialize_eventloops_for_actor_concurrency_group(
  155. self, const c_vector[CConcurrencyGroup] &c_defined_concurrency_groups)
  156. cdef python_scheduling_strategy_to_c(
  157. self, python_scheduling_strategy,
  158. CSchedulingStrategy *c_scheduling_strategy)
  159. cdef python_label_match_expressions_to_c(
  160. self, python_expressions,
  161. CLabelMatchExpressions *c_expressions)
  162. cdef CObjectID allocate_dynamic_return_id_for_generator(
  163. self,
  164. const CAddress &owner_address,
  165. const CTaskID &task_id,
  166. return_size,
  167. generator_index,
  168. is_async_actor)
  169. cdef class FunctionDescriptor:
  170. cdef:
  171. CFunctionDescriptor descriptor