easy_ml 0.2.0.pre.rc77 → 0.2.0.pre.rc78

data/app/models/easy_ml/model.rb

@@ -45,7 +45,7 @@ module EasyML
     MODEL_NAMES = MODEL_OPTIONS.keys.freeze
     MODEL_CONSTANTS = MODEL_OPTIONS.values.map(&:constantize)
 
-    add_configuration_attributes :task, :objective, :hyperparameters, :callbacks, :metrics
+    add_configuration_attributes :task, :objective, :hyperparameters, :callbacks, :metrics, :weights_column
     MODEL_CONSTANTS.flat_map(&:configuration_attributes).each do |attribute|
       add_configuration_attributes attribute
     end
@@ -179,6 +179,8 @@ module EasyML
     end
 
     def actually_train(&progress_block)
+      raise untrainable_error unless trainable?
+
       lock_model do
         run = pending_run
         run.wrap_training do
@@ -258,7 +260,7 @@ module EasyML
 
     def formatted_version
       return nil unless version
-      Time.strptime(version, "%Y%m%d%H%M%S").strftime("%B %-d, %Y at %-l:%M %p")
+      UTC.parse(version).in_time_zone(EasyML::Configuration.timezone).strftime("%B %-d, %Y at %-l:%M %p")
     end
 
     def last_run_at
@@ -277,6 +279,22 @@ module EasyML
     alias_method :latest_version, :inference_version
     alias_method :deployed, :inference_version
 
+    def trainable?
+      adapter.trainable?
+    end
+
+    def untrainable_columns
+      adapter.untrainable_columns
+    end
+
+    def untrainable_error
+      %Q(
+        Cannot train dataset containing null values!
+        Apply preprocessing to the following columns:
+        #{untrainable_columns.join(", ")}
+      )
+    end
+
     def predict(xs)
       load_model!
       unless xs.is_a?(XGBoost::DMatrix)

data/app/models/easy_ml/models/xgboost/evals_callback.rb

@@ -36,7 +36,7 @@ module EasyML
           if tuner.present?
             [tuner.x_valid, tuner.y_valid]
           else
-            model.dataset.valid(split_ys: true)
+            model.dataset.valid(split_ys: true, lazy: true)
           end
         end
 
@@ -47,7 +47,8 @@ module EasyML
           if epoch % log_frequency == 0
             model.adapter.external_model = booster
             x_valid, y_valid = valid_dataset
-            @preprocessed ||= model.preprocess(x_valid)
+            x_valid = x_valid.select(model.dataset.col_order(inference: true))
+            @preprocessed ||= model.preprocess(x_valid, y_valid)
             y_pred = model.predict(@preprocessed)
             dataset = model.dataset.valid(all_columns: true)
 

data/app/models/easy_ml/models/xgboost.rb

@@ -421,11 +421,11 @@ module EasyML
       def prepare_data
         if @d_train.nil?
           col_order = dataset.col_order
-          x_sample, y_sample = dataset.train(split_ys: true, limit: 5, select: col_order)
+          x_sample, y_sample = dataset.train(split_ys: true, limit: 5, select: col_order, lazy: true)
           preprocess(x_sample, y_sample) # Ensure we fail fast if the dataset is misconfigured
-          x_train, y_train = dataset.train(split_ys: true, select: col_order)
-          x_valid, y_valid = dataset.valid(split_ys: true, select: col_order)
-          x_test, y_test = dataset.test(split_ys: true, select: col_order)
+          x_train, y_train = dataset.train(split_ys: true, select: col_order, lazy: true)
+          x_valid, y_valid = dataset.valid(split_ys: true, select: col_order, lazy: true)
+          x_test, y_test = dataset.test(split_ys: true, select: col_order, lazy: true)
           @d_train = preprocess(x_train, y_train)
           @d_valid = preprocess(x_valid, y_valid)
           @d_test = preprocess(x_test, y_test)
@@ -434,21 +434,60 @@ module EasyML
         [@d_train, @d_valid, @d_test]
       end
 
+      def trainable?
+        untrainable_columns.empty?
+      end
+
+      def untrainable_columns
+        df = model.dataset.processed.data(lazy: true)
+
+        columns = df.columns
+        selects = columns.map do |col|
+          Polars.col(col).null_count.alias(col)
+        end
+        null_info = df.select(selects).collect
+        null_info.to_hashes.first.compact
+        col_list = null_info.to_hashes.first.transform_values { |v| v > 0 ? v : nil }.compact.keys
+
+        model.dataset.regular_columns(col_list)
+      end
+
       def preprocess(xs, ys = nil)
         return xs if xs.is_a?(::XGBoost::DMatrix)
+        weights_col = model.weights_column || nil
+
+        if weights_col == model.dataset.target
+          raise ArgumentError, "Weight column cannot be the target column"
+        end
+
+        # Extract feature columns (all columns except label and weight)
+        feature_cols = xs.columns
+        feature_cols -= [weights_col] if weights_col
+        lazy = xs.is_a?(Polars::LazyFrame)
+
+        # Get features, labels and weights
+        features = lazy ? xs.select(feature_cols).collect.to_numo : xs.select(feature_cols).to_numo
+        weights = weights_col ? (lazy ? xs.select(weights_col).collect.to_numo : xs.select(weights_col).to_numo) : nil
+        weights = weights.flatten if weights
+        if ys.present?
+          ys = ys.is_a?(Array) ? Polars::Series.new(ys) : ys
+          labels = lazy ? ys.collect.to_numo.flatten : ys.to_numo.flatten
+        else
+          labels = nil
+        end
+
+        kwargs = {
+          label: labels,
+          weight: weights,
+        }.compact
 
-        orig_xs = xs.dup
-        column_names = xs.columns
-        xs = _preprocess(xs)
-        ys = ys.nil? ? nil : _preprocess(ys).flatten
-        kwargs = { label: ys }.compact
         begin
-          ::XGBoost::DMatrix.new(xs, **kwargs).tap do |dmat|
-            dmat.feature_names = column_names
+          ::XGBoost::DMatrix.new(features, **kwargs).tap do |dmatrix|
+            dmatrix.feature_names = feature_cols
           end
         rescue StandardError => e
-          problematic_columns = orig_xs.schema.select { |k, v| [Polars::Categorical, Polars::String].include?(v) }
-          problematic_xs = orig_xs.select(problematic_columns.keys)
+          problematic_columns = xs.schema.select { |k, v| [Polars::Categorical, Polars::String].include?(v) }
+          problematic_xs = lazy ? xs.lazy.select(problematic_columns.keys).collect : xs.select(problematic_columns.keys)
           raise %(
             Error building data for XGBoost.
             Apply preprocessing to columns 
@@ -501,29 +540,6 @@ module EasyML
         cb_container.after_iteration(@booster, current_iteration, d_train, evals)
       end
 
-      def _preprocess(df)
-        return df if df.is_a?(Array)
-
-        df.to_a.map do |row|
-          row.values.map do |value|
-            case value
-            when Time
-              value.to_i # Convert Time to Unix timestamp
-            when Date
-              value.to_time.to_i # Convert Date to Unix timestamp
-            when String
-              value
-            when TrueClass, FalseClass
-              value ? 1.0 : 0.0 # Convert booleans to 1.0 and 0.0
-            when Integer
-              value
-            else
-              value.to_f # Ensure everything else is converted to a float
-            end
-          end
-        end
-      end
-
       def initialize_model
         @xgboost_model = model_class.new(n_estimators: @hyperparameters.to_h.dig(:n_estimators))
         if block_given?

data/app/models/easy_ml/retraining_run.rb

@@ -150,7 +150,7 @@ module EasyML
 
       training_model.dataset.refresh
       evaluator = retraining_job.evaluator.symbolize_keys
-      x_test, y_test = training_model.dataset.test(split_ys: true)
+      x_test, y_test = training_model.dataset.test(split_ys: true, all_columns: true)
       y_pred = training_model.predict(x_test)
 
       metric = evaluator[:metric].to_sym

data/app/serializers/easy_ml/model_serializer.rb

@@ -27,6 +27,7 @@ module EasyML
                :model_type,
                :task,
                :objective,
+               :weights_column,
                :metrics,
                :dataset_id,
                :status,

data/lib/easy_ml/core/tuner.rb

@@ -8,7 +8,7 @@ module EasyML
                     :metrics, :objective, :n_trials, :direction, :evaluator,
                     :study, :results, :adapter, :tune_started_at, :x_valid, :y_valid,
                     :project_name, :job, :current_run, :trial_enumerator, :progress_block,
-                    :tuner_job, :dataset
+                    :tuner_job, :dataset, :x_normalized
 
       def initialize(options = {})
         @model = options[:model]
@@ -73,9 +73,12 @@ module EasyML
         model.task = task
 
         model.dataset.refresh if model.dataset.needs_refresh?
-        x_valid, y_valid = model.dataset.valid(split_ys: true, select: model.dataset.col_order)
+        x_valid, y_valid = model.dataset.valid(split_ys: true, all_columns: true)
+        x_normalized = model.dataset.normalize(x_valid, inference: true)
+        x_normalized = model.preprocess(x_normalized)
         self.x_valid = x_valid
         self.y_valid = y_valid
+        self.x_normalized = x_normalized
         self.dataset = model.dataset.valid(all_columns: true)
         adapter.tune_started_at = tune_started_at
         adapter.x_valid = x_valid
@@ -99,7 +102,7 @@ module EasyML
             @study.tell(@current_trial, result)
           rescue StandardError => e
             puts EasyML::Event.easy_ml_context(e.backtrace)
-            @tuner_run.update!(status: :failed, hyperparameters: {})
+            @tuner_run.update!(status: :failed, hyperparameters: model.hyperparameters.to_h)
             puts "Optuna failed with: #{e.message}"
             raise e
           end
@@ -138,7 +141,7 @@ module EasyML
           end
         end
 
-        y_pred = model.predict(x_valid)
+        y_pred = model.predict(x_normalized)
         model.metrics = metrics
         metrics = model.evaluate(y_pred: y_pred, y_true: y_valid, x_true: x_valid, dataset: dataset)
         metric = metrics.symbolize_keys.dig(model.evaluator[:metric].to_sym)

data/lib/easy_ml/data/dataset_manager/writer/base.rb

@@ -36,6 +36,10 @@ module EasyML
             clear_unique_id
           end
 
+          def unlock!
+            clear_all_keys
+          end
+
           private
 
           def files
@@ -50,6 +54,10 @@ module EasyML
             safe_write(df, unique_path(subdir: subdir))
           end
 
+          def acquire_lock(key, &block)
+            Support::Lockable.with_lock("#{key}:lock", wait_timeout: 2, &block)
+          end
+
           def unique_path(subdir: nil)
             filename = [filenames, unique_id(subdir: subdir), "parquet"].compact.join(".")
 
@@ -63,15 +71,18 @@ module EasyML
           end
 
           def clear_all_keys
-            keys = list_keys
-            Support::Lockable.with_lock(keys, wait_timeout: 2) do |suo|
-              suo.client.del(keys)
+            list_keys.each { |key| unlock_file(key) }
+          end
+
+          def unlock_file(key)
+            acquire_lock(key) do |suo|
+              suo.client.del(key)
             end
           end
 
           def clear_unique_id(subdir: nil)
             key = unique_id_key(subdir: subdir)
-            Support::Lockable.with_lock(key, wait_timeout: 2) do |suo|
+            acquire_lock(key) do |suo|
               suo.client.del(key)
             end
           end
@@ -83,7 +94,7 @@ module EasyML
           def add_key(key)
             keylist = unique_id_key(subdir: "keylist")
 
-            Support::Lockable.with_lock(keylist, wait_timeout: 2) do |suo|
+            acquire_lock(keylist) do |suo|
               suo.client.sadd(keylist, key)
             end
           end
@@ -91,14 +102,20 @@ module EasyML
           def list_keys
             keylist = unique_id_key(subdir: "keylist")
 
-            Support::Lockable.with_lock(keylist, wait_timeout: 2) do |suo|
-              suo.client.smembers(keylist)
+            acquire_lock(keylist) do |suo|
+              if suo.client.type(keylist) == "set"
+                suo.client.smembers(keylist)
+              else
+                suo.client.del(keylist)
+                []
+              end
             end
           end
 
           def key_exists?(key)
             keylist = unique_id_key(subdir: "keylist")
-            Support::Lockable.with_lock(keylist, wait_timeout: 2) do |suo|
+
+            acquire_lock(keylist) do |suo|
               suo.client.sismember(keylist, key)
             end
           end
@@ -107,7 +124,7 @@ module EasyML
             key = unique_id_key(subdir: subdir)
             add_key(key)
 
-            Support::Lockable.with_lock(key, wait_timeout: 2) do |suo|
+            acquire_lock(key) do |suo|
               redis = suo.client
 
               seq = (redis.get(key) || "0").to_i

data/lib/easy_ml/data/dataset_manager/writer.rb

@@ -15,7 +15,7 @@ module EasyML
         ]
 
         attr_accessor :filenames, :root_dir, :partition,
-                      :append_only, :primary_key, :options
+                      :primary_key, :options, :append_only, :named
 
         def initialize(options)
           @root_dir = options.dig(:root_dir)
@@ -27,6 +27,10 @@ module EasyML
           @options = options
         end
 
+        def unlock!
+          adapter_class.new(options).unlock!
+        end
+
         def store(df, *args)
           adapter_class.new(options.merge!(df: df)).store(*args)
         end

data/lib/easy_ml/data/dataset_manager.rb

@@ -67,6 +67,8 @@ module EasyML
         Reader.sha(root_dir)
       end
 
+      # Transform CSV files into Parquet files, of all the same datatype.
+      # Learn datatypes of columns and store schema.
       def normalize
         Normalizer.normalize(root_dir)
       end
@@ -75,14 +77,18 @@ module EasyML
         query
       end
 
-      def store(df, *args)
-        writer.store(df, *args)
+      def unlock!
+        writer.unlock!
       end
 
       def compact
         writer.compact
       end
 
+      def store(df, *args)
+        writer.store(df, *args)
+      end
+
       def cp(from, to)
         writer.cp(from, to)
       end

data/lib/easy_ml/data/polars_column.rb

@@ -124,11 +124,25 @@ module EasyML
       # @param series [Polars::Series] The string series to analyze
       # @return [Symbol] One of :datetime, :text, or :categorical
       def determine_string_type(series)
-        if EasyML::Data::DateConverter.maybe_convert_date(Polars::DataFrame.new({ temp: series }),
-                                                          :temp)[:temp].dtype.is_a?(Polars::Datetime)
-          :datetime
-        else
-          categorical_or_text?(series)
+        # Try to parse as numeric first
+        begin
+          # Try integer first
+          series.cast(Polars::Int64)
+          return :numeric
+        rescue StandardError
+          begin
+            # Try float if integer fails
+            series.cast(Polars::Float64)
+            return :numeric
+          rescue StandardError
+            # If not numeric, check for datetime or categorical
+            if EasyML::Data::DateConverter.maybe_convert_date(Polars::DataFrame.new({ temp: series }),
+                                                              :temp)[:temp].dtype.is_a?(Polars::Datetime)
+              :datetime
+            else
+              categorical_or_text?(series)
+            end
+          end
         end
       end
 

data/lib/easy_ml/engine.rb

@@ -55,7 +55,7 @@ module EasyML
       Polars.enable_string_cache
     end
 
-    if %w[db:migrate db:migrate:status db:setup db:drop assets:precompile].include?(ARGV.first)
+    if %w[db:create db:migrate db:migrate:status db:setup db:drop assets:precompile].include?(ARGV.first)
       config.eager_load_paths = config.eager_load_paths.without(config.eager_load_paths.map(&:to_s).grep(/easy_ml/).map { |p| Pathname.new(p) })
     else
       config.after_initialize do
@@ -77,6 +77,21 @@ module EasyML
       end
     end
 
+    unless %w[db:create db:migrate db:migrate:status db:setup db:drop assets:precompile].include?(ARGV.first)
+      initializer "easy_ml.configure_secrets" do
+        EasyML::Configuration.configure do |config|
+          raise "S3_ACCESS_KEY_ID is missing. Set ENV['S3_ACCESS_KEY_ID']" unless ENV["S3_ACCESS_KEY_ID"]
+          raise "S3_SECRET_ACCESS_KEY is missing. Set ENV['S3_SECRET_ACCESS_KEY']" unless ENV["S3_SECRET_ACCESS_KEY"]
+
+          config.s3_access_key_id = ENV["S3_ACCESS_KEY_ID"]
+          config.s3_secret_access_key = ENV["S3_SECRET_ACCESS_KEY"]
+          config.s3_region = ENV["S3_REGION"] ? ENV["S3_REGION"] : "us-east-1"
+          config.timezone = ENV["TIMEZONE"].present? ? ENV["TIMEZONE"] : "America/New_York"
+          config.wandb_api_key = ENV["WANDB_API_KEY"] if ENV["WANDB_API_KEY"]
+        end
+      end
+    end
+
     initializer "easy_ml.check_pending_migrations" do
       if defined?(Rails::Server)
         config.after_initialize do
@@ -96,19 +111,6 @@ module EasyML
       end
     end
 
-    initializer "easy_ml.configure_secrets" do
-      EasyML::Configuration.configure do |config|
-        raise "S3_ACCESS_KEY_ID is missing. Set ENV['S3_ACCESS_KEY_ID']" unless ENV["S3_ACCESS_KEY_ID"]
-        raise "S3_SECRET_ACCESS_KEY is missing. Set ENV['S3_SECRET_ACCESS_KEY']" unless ENV["S3_SECRET_ACCESS_KEY"]
-
-        config.s3_access_key_id = ENV["S3_ACCESS_KEY_ID"]
-        config.s3_secret_access_key = ENV["S3_SECRET_ACCESS_KEY"]
-        config.s3_region = ENV["S3_REGION"] if ENV["S3_REGION"]
-        config.timezone = ENV["TIMEZONE"].present? ? ENV["TIMEZONE"] : "America/New_York"
-        config.wandb_api_key = ENV["WANDB_API_KEY"] if ENV["WANDB_API_KEY"]
-      end
-    end
-
     initializer "easy_ml.setup_generators" do |app|
       generators_path = EasyML::Engine.root.join("lib/easy_ml/railtie/generators")
       generators_dirs = Dir[File.join(generators_path, "**", "*.rb")]

data/lib/easy_ml/feature_store.rb

@@ -5,19 +5,22 @@ module EasyML
     def initialize(feature)
       @feature = feature
 
-      datasource_config = feature.dataset.datasource.configuration || {}
-
-      options = {
-        root_dir: feature_dir,
-        filenames: "feature",
-        append_only: false,
-        primary_key: feature.primary_key&.first,
-        partition_size: batch_size,
-        s3_bucket: datasource_config.dig("s3_bucket") || EasyML::Configuration.s3_bucket,
-        s3_prefix: s3_prefix,
-        polars_args: datasource_config.dig("polars_args"),
-      }.compact
-      super(options)
+      datasource_config = feature&.dataset&.datasource&.configuration
+      if datasource_config
+        options = {
+          root_dir: feature_dir,
+          filenames: "feature",
+          append_only: false,
+          primary_key: feature.primary_key&.first,
+          partition_size: batch_size,
+          s3_bucket: datasource_config.dig("s3_bucket") || EasyML::Configuration.s3_bucket,
+          s3_prefix: s3_prefix,
+          polars_args: datasource_config.dig("polars_args"),
+        }.compact
+        super(options)
+      else
+        super({ root_dir: "" })
+      end
     end
 
     def cp(old_version, new_version)
@@ -30,7 +33,7 @@ module EasyML
       files_to_cp = Dir.glob(Pathname.new(old_dir).join("**/*")).select { |f| File.file?(f) }
 
       files_to_cp.each do |file|
-        target_file = file.gsub(old_version.to_s, new_version.to_s)
+        target_file = file.gsub(old_dir, new_dir)
         FileUtils.mkdir_p(File.dirname(target_file))
         FileUtils.cp(file, target_file)
       end
@@ -46,9 +49,9 @@ module EasyML
       File.join(
         Rails.root,
         "easy_ml/datasets",
-        feature.dataset.name.parameterize.gsub("-", "_"),
+        feature&.dataset&.name&.parameterize&.gsub("-", "_"),
         "features",
-        feature.name.parameterize.gsub("-", "_"),
+        feature&.name&.parameterize&.gsub("-", "_"),
         version.to_s
       )
     end

data/lib/easy_ml/support/lockable.rb

@@ -22,7 +22,7 @@ module EasyML
       end
 
       def self.lock_client(key, wait_timeout: 0.1, stale_timeout: 60 * 10, resources: 1)
-        Suo::Client::Redis.new(prefixed_key(key), {
+        Suo::Client::Redis.new(key, {
           acquisition_timeout: wait_timeout,
           stale_lock_expiry: stale_timeout,
           resources: resources,
@@ -30,10 +30,6 @@ module EasyML
         })
       end
 
-      def self.prefixed_key(key)
-        "easy_ml:#{key}"
-      end
-
       # Execute a block with a Redis lock
       def self.with_lock(key, wait_timeout: 0.1, stale_timeout: 60 * 10, resources: 1)
         lock_key = nil

data/lib/easy_ml/version.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module EasyML
-  VERSION = "0.2.0-rc77"
+  VERSION = "0.2.0-rc78"
 
   module Version
   end

data/public/easy_ml/assets/.vite/manifest.json

@@ -1,6 +1,6 @@
 {
   "entrypoints/Application.tsx": {
-    "file": "assets/entrypoints/Application.tsx-B1qLZuyu.js",
+    "file": "assets/entrypoints/Application.tsx-Bbf3mD_b.js",
     "name": "entrypoints/Application.tsx",
     "src": "entrypoints/Application.tsx",
     "isEntry": true,