static_embeddings 1.5.7 → 1.5.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: a8e174dd93ba3b3e9ee25047fd28ea5ba980d096a4922c0b8325c366208bff65
-  data.tar.gz: 741e0cc35791d9a335c82a05aed18e3388232828824168ff59ed1751c281e8a8
+  metadata.gz: cbb2b16d03e01ba91c8c856890e98537ff81c4b4961a4d4172dc2e8a27fb46e3
+  data.tar.gz: d182451276da091e5410fa5dcbf93ec08e773bfa51cc6226f0c244db132e3656
 SHA512:
-  metadata.gz: 58462ed4c203fb657630d092e1a4aa85857e5d44106bb97feae26f12579fc4b6d88d20632e462624ad0cb93e1d0d0e735f91ee12aaec80b637d29b313e6dc53d
-  data.tar.gz: 80e9897969dfe6aace3d851f1a11ffb6610b628cddc35239301a676b2ce32d97d1a8a614ae6dca0526703e61d6d813af89965d0906c593392e66cdd25f7e16da
+  metadata.gz: 54d96917d4a2fd738787ee12bd57a63f3793893a4abe3e59a9b860efe8508a895cb4d1cbd865f4f5692af1b2aa46eea21b17b55ed090b98a9af8df59946feadc
+  data.tar.gz: da73a4ead8792691b1e8433ecb7487f121e1dda88c61054e56f0f05b56c43f922a49643f1668da4c34be1258f654f134975cae5db875075f2a99e9bd2ea5256a

data/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 1.5.8
+
+### Fixed
+
+- Cancellation flags shared by Ruby's unblock callback and native worker threads now use
+  atomic loads and stores, removing a C data race during interrupted embedding, tokenization,
+  and top-k operations.
+
 ## 1.5.7
 
 ### Fixed

data/ext/static_embeddings/se_embed.c

@@ -119,7 +119,7 @@ void se_l2_normalize(float *vec, uint32_t dim) {
 
 static se_status_t embed_ids_core(const se_model_t *model, se_scratch_t *sc, const uint32_t *ids,
                                   size_t n, float *out, se_error_t *err,
-                                  volatile sig_atomic_t *cancelled) {
+                                  se_cancel_flag_t *cancelled) {
     const uint32_t dim = model->meta.dim;
     float *acc = sc->acc;
     memset(acc, 0, (size_t)dim * sizeof(float));
@@ -129,7 +129,7 @@ static se_status_t embed_ids_core(const se_model_t *model, se_scratch_t *sc, con
 
     size_t used = 0;
     for (size_t i = 0; i < n; i++) {
-        if (((i & 255u) == 0) && cancelled && *cancelled) {
+        if (((i & 255u) == 0) && se_cancel_requested(cancelled)) {
             se_error_set(err, SE_ERR_INTERNAL, "operation cancelled");
             return SE_ERR_INTERNAL;
         }
@@ -179,7 +179,7 @@ static se_status_t embed_ids_core(const se_model_t *model, se_scratch_t *sc, con
 
 se_status_t se_embed_ids(const se_model_t *model, se_scratch_t *sc, const uint32_t *ids,
                          size_t n_ids, float *out, se_token_stats_t *stats, se_error_t *err,
-                         volatile sig_atomic_t *cancelled) {
+                         se_cancel_flag_t *cancelled) {
     if (n_ids > UINT32_MAX) {
         se_error_set(err, SE_ERR_OOM, "too many token ids");
         return SE_ERR_OOM;
@@ -196,7 +196,7 @@ se_status_t se_embed_ids(const se_model_t *model, se_scratch_t *sc, const uint32
 
 se_status_t se_embed_one(const se_model_t *model, se_scratch_t *sc, const uint8_t *input,
                          size_t input_len, uint32_t max_tokens, float *out, se_token_stats_t *stats,
-                         se_error_t *err, volatile sig_atomic_t *cancelled) {
+                         se_error_t *err, se_cancel_flag_t *cancelled) {
     se_status_t rc = se_tokenize(model, sc, input, input_len, max_tokens, SE_TOKEN_LIMIT_USABLE,
                                  stats, err, cancelled);
     if (rc != SE_OK)

data/ext/static_embeddings/se_internal.h

@@ -3,9 +3,18 @@
 
 #include <stddef.h>
 #include <stdint.h>
-#include <signal.h>
 #include <limits.h>
 
+typedef int se_cancel_flag_t;
+
+static inline int se_cancel_requested(const se_cancel_flag_t *flag) {
+    return flag && __atomic_load_n(flag, __ATOMIC_RELAXED);
+}
+
+static inline void se_cancel_request(se_cancel_flag_t *flag) {
+    __atomic_store_n(flag, 1, __ATOMIC_RELAXED);
+}
+
 #define SE_MAGIC          "SEMBv1\0\0"
 #define SE_MAGIC_LEN      8
 #define SE_FORMAT_VERSION 3u
@@ -277,7 +286,7 @@ typedef struct {
     float *best_score;
     int cosine;
     float inv_query_norm;
-    volatile sig_atomic_t cancelled;
+    se_cancel_flag_t cancelled;
 } se_topk_job_t;
 
 SE_STATIC_ASSERT(sizeof(se_vocab_slot_t) == 16, vocab_slot_size);
@@ -415,14 +424,14 @@ typedef enum { SE_TOKEN_LIMIT_RAW = 0, SE_TOKEN_LIMIT_USABLE = 1 } se_token_limi
 
 se_status_t se_tokenize(const se_model_t *model, se_scratch_t *scratch, const uint8_t *input,
                         size_t input_len, uint32_t max_tokens, se_token_limit_t limit_mode,
-                        se_token_stats_t *stats, se_error_t *err, volatile sig_atomic_t *cancelled);
+                        se_token_stats_t *stats, se_error_t *err, se_cancel_flag_t *cancelled);
 
 se_status_t se_embed_one(const se_model_t *model, se_scratch_t *scratch, const uint8_t *input,
                          size_t input_len, uint32_t max_tokens, float *out, se_token_stats_t *stats,
-                         se_error_t *err, volatile sig_atomic_t *cancelled);
+                         se_error_t *err, se_cancel_flag_t *cancelled);
 se_status_t se_embed_ids(const se_model_t *model, se_scratch_t *scratch, const uint32_t *ids,
                          size_t n_ids, float *out, se_token_stats_t *stats, se_error_t *err,
-                         volatile sig_atomic_t *cancelled);
+                         se_cancel_flag_t *cancelled);
 
 void se_l2_normalize(float *vec, uint32_t dim);
 size_t se_model_warmup(const se_model_t *model);

data/ext/static_embeddings/se_tokenizer.c

@@ -495,11 +495,11 @@ typedef struct {
     int segment_ascii;
     se_token_stats_t *stats;
     se_error_t *err;
-    volatile sig_atomic_t *cancelled;
+    se_cancel_flag_t *cancelled;
 } token_state_t;
 
 static int token_cancelled(const token_state_t *st) {
-    return st->cancelled && *st->cancelled;
+    return se_cancel_requested(st->cancelled);
 }
 
 static se_status_t oom(token_state_t *st, const char *where) {
@@ -954,7 +954,7 @@ static se_status_t tokenize_ascii_run(token_state_t *st, const uint8_t *input, s
 se_status_t se_tokenize(const se_model_t *model, se_scratch_t *sc, const uint8_t *input,
                         size_t input_len, uint32_t max_tokens, se_token_limit_t limit_mode,
                         se_token_stats_t *stats, se_error_t *err,
-                        volatile sig_atomic_t *cancelled) {
+                        se_cancel_flag_t *cancelled) {
     memset(stats, 0, sizeof(*stats));
 
     token_state_t st;

data/ext/static_embeddings/se_topk.c

@@ -195,7 +195,7 @@ static float cosine_score(float dot, float row_sq, float inv_query_norm) {
 void *se_topk_execute(void *arg) {
     se_topk_job_t *job = (se_topk_job_t *)arg;
     for (size_t r = 0; r < job->rows; r++) {
-        if ((r & 1023u) == 0 && job->cancelled)
+        if ((r & 1023u) == 0 && se_cancel_requested(&job->cancelled))
             return NULL;
         float score;
 

data/ext/static_embeddings/static_embeddings.c

@@ -9,7 +9,6 @@
 #include <float.h>
 #include <limits.h>
 #include <math.h>
-#include <signal.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -226,7 +225,7 @@ typedef struct {
     size_t next_index;
     int failed;
     size_t failed_index;
-    volatile sig_atomic_t cancelled;
+    se_cancel_flag_t cancelled;
     se_error_t error;
 } batch_job_t;
 
@@ -244,7 +243,7 @@ static void batch_worker_run(batch_job_t *job, se_scratch_t *scratch) {
     const uint32_t dim = job->model->meta.dim;
 
     for (;;) {
-        if (job->cancelled) {
+        if (se_cancel_requested(&job->cancelled)) {
             se_error_set(&job->error, SE_ERR_INTERNAL, "operation cancelled");
             job->failed = 1;
             return;
@@ -293,7 +292,7 @@ static void *batch_execute(void *arg) {
 static void unblock_cancel(void *arg) {
     batch_job_t *job = (batch_job_t *)arg;
     if (job)
-        job->cancelled = 1;
+        se_cancel_request(&job->cancelled);
 }
 
 static VALUE current_fiber_scheduler(void) {
@@ -358,7 +357,7 @@ static VALUE fiber_worker_ensure(VALUE wrapper) {
     fiber_worker_t *w;
     TypedData_Get_Struct(wrapper, fiber_worker_t, &fiber_worker_type, w);
     if (!NIL_P(w->thread)) {
-        w->job->cancelled = 1;
+        se_cancel_request(&w->job->cancelled);
         rb_funcall(w->thread, id_kill, 0);
         rb_funcall(w->thread, id_join, 0);
     }
@@ -885,7 +884,7 @@ static VALUE embed_texts_internal(VALUE self, VALUE texts, int is_array, size_t
         int state = 0;
         rb_protect(protected_run_job, (VALUE)(uintptr_t)&job, &state);
         if (state) {
-            job.cancelled = 1;
+            se_cancel_request(&job.cancelled);
             free_input(&job.input);
             embed_run_free(&run);
             rb_jump_tag(state);
@@ -1151,7 +1150,7 @@ typedef struct {
     size_t n_ids;
     float *out;
     se_token_stats_t stats;
-    volatile sig_atomic_t cancelled;
+    se_cancel_flag_t cancelled;
     se_error_t error;
 } ids_job_t;
 
@@ -1200,7 +1199,7 @@ static void *ids_execute(void *arg) {
 static void ids_unblock_cancel(void *arg) {
     ids_job_t *job = (ids_job_t *)arg;
     if (job)
-        job->cancelled = 1;
+        se_cancel_request(&job->cancelled);
 }
 
 static VALUE embed_token_ids_body(VALUE arg) {
@@ -1246,7 +1245,7 @@ static VALUE embed_token_ids_body(VALUE arg) {
         run->owner->active_ops--;
         run->active = 0;
         rb_thread_check_ints();
-        if (job.cancelled)
+        if (se_cancel_requested(&job.cancelled))
             rb_raise(rb_eInterrupt, "operation cancelled");
     } else {
         ids_execute(&job);
@@ -1396,7 +1395,7 @@ static int ptr_is_float_aligned(const void *ptr) {
 static void topk_unblock_cancel(void *arg) {
     se_topk_job_t *job = (se_topk_job_t *)arg;
     if (job)
-        job->cancelled = 1;
+        se_cancel_request(&job->cancelled);
 }
 
 static VALUE topk_body(VALUE arg) {
@@ -1425,7 +1424,7 @@ static VALUE topk_body(VALUE arg) {
         rb_thread_check_ints();
     }
 
-    if (run->job.cancelled)
+    if (se_cancel_requested(&run->job.cancelled))
         rb_raise(rb_eInterrupt, "operation cancelled");
 
     VALUE result = rb_ary_new_capa(run->job.k);

data/lib/static_embeddings/version.rb

@@ -1,3 +1,3 @@
 module StaticEmbeddings
-  VERSION = "1.5.7"
+  VERSION = "1.5.8"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: static_embeddings
 version: !ruby/object:Gem::Version
-  version: 1.5.7
+  version: 1.5.8
 platform: ruby
 authors:
 - Roman Haydarov