easy_ml 0.2.0.pre.rc40 → 0.2.0.pre.rc43

data/app/models/easy_ml/feature.rb

@@ -17,6 +17,7 @@
 #  refresh_every    :bigint
 #  created_at       :datetime         not null
 #  updated_at       :datetime         not null
+#  workflow_status  :string
 #
 module EasyML
   class Feature < ActiveRecord::Base
@@ -24,6 +25,11 @@ module EasyML
     include Historiographer::Silent
     historiographer_mode :snapshot_only
 
+    enum workflow_status: {
+      analyzing: "analyzing",
+      ready: "ready",
+      failed: "failed",
+    }
     class << self
       def compute_sha(feature_class)
         require "digest"
@@ -135,13 +141,22 @@ module EasyML
       adapter.respond_to?(:batch) || config.dig(:batch_size).present?
     end
 
+    def primary_key
+      pkey = config.dig(:primary_key)
+      if pkey.is_a?(Array)
+        pkey
+      else
+        [pkey]
+      end
+    end
+
     def numeric_primary_key?
       if primary_key.nil?
         return false unless should_be_batchable?
         raise "Couldn't find primary key for feature #{feature_class}, check your feature class"
       end
 
-      dataset.raw.data(limit: 1, select: primary_key)[primary_key].to_a.flat_map(&:values).all? do |value|
+      dataset.raw.data(limit: 1, select: primary_key)[primary_key].to_a.flat_map { |h| h.respond_to?(:values) ? h.values : h }.all? do |value|
         case value
         when String then value.match?(/\A[-+]?\d+(\.\d+)?\z/)
         else
@@ -171,22 +186,25 @@ module EasyML
           unless primary_key.present?
             raise "Couldn't find primary key for feature #{feature_class}, check your feature class"
           end
-          df = reader.query(select: [primary_key.first])
+          df = reader.query(select: primary_key)
         rescue => e
           raise "Couldn't find primary key #{primary_key.first} for feature #{feature_class}: #{e.message}"
         end
         return [] if df.nil?
 
         min_id = df[primary_key.first].min
-        max_id = df[primary_key.first].max
+        max_id = df[primary_key.last].max
       end
 
-      (min_id..max_id).step(batch_size).map do |batch_start|
+      (min_id..max_id).step(batch_size).map.with_index do |batch_start, idx|
         batch_end = [batch_start + batch_size, max_id + 1].min - 1
         {
           feature_id: id,
           batch_start: batch_start,
           batch_end: batch_end,
+          batch_number: feature_position,
+          subbatch_number: idx,
+          parent_batch_id: Random.uuid,
         }
       end
     end
@@ -196,13 +214,16 @@ module EasyML
     end
 
     def fit(features: [self], async: false)
-      jobs = features.flat_map(&:build_batches)
+      ordered_features = features.sort_by(&:feature_position)
+      jobs = ordered_features.map(&:build_batches)
+
       if async
-        EasyML::ComputeFeatureJob.enqueue_batch(jobs)
+        EasyML::ComputeFeatureJob.enqueue_ordered_batches(jobs)
       else
-        jobs.each do |job|
+        jobs.flatten.each do |job|
           EasyML::ComputeFeatureJob.perform(nil, job)
         end
+        features.each(&:after_fit) unless features.any?(&:failed?)
       end
     end
 
@@ -266,13 +287,11 @@ module EasyML
           batch_df = adapter.fit(df, self, batch_args)
         end
       end
-      raise "Feature #{feature_class}#fit must return a dataframe" unless batch_df.present?
-      store(batch_df)
-      updates = {
-        applied_at: Time.current,
-        needs_fit: false,
-      }.compact
-      update!(updates)
+      if batch_df.present?
+        store(batch_df)
+      else
+        "Feature #{feature_class}#fit should return a dataframe, received #{batch_df.class}"
+      end
       batch_df
     end
 
@@ -335,6 +354,7 @@ module EasyML
     def apply_defaults
       self.name ||= self.feature_class.demodulize.titleize
       self.version ||= 1
+      self.workflow_status ||= :ready
     end
 
     def needs_columns
@@ -371,6 +391,15 @@ module EasyML
         (should_be_batchable? ? 10_000 : nil)
     end
 
+    def after_fit
+      updates = {
+        applied_at: Time.current,
+        needs_fit: false,
+        workflow_status: :ready,
+      }.compact
+      update!(updates)
+    end
+
     private
 
     def bulk_update_positions(features)

data/app/models/easy_ml/model.rb

@@ -250,6 +250,7 @@ module EasyML
       bump_version(force: true)
       path = model_file.full_path(version)
       full_path = adapter.save_model_file(path)
+      puts "saving model to #{full_path}"
       model_file.upload(full_path)
 
       model_file.save
@@ -266,6 +267,7 @@ module EasyML
     end
 
     def cleanup
+      puts "keeping files #{files_to_keep}"
       get_model_file&.cleanup(files_to_keep)
     end
 
@@ -488,13 +490,9 @@ module EasyML
     end
 
     def root_dir
-      persisted = read_attribute(:root_dir)
+      relative_dir = read_attribute(:root_dir) || default_root_dir
 
-      if persisted.present? && !persisted.blank?
-        EasyML::Engine.root_dir.join(persisted).to_s
-      else
-        default_root_dir
-      end
+      EasyML::Engine.root_dir.join(relative_dir).to_s
     end
 
     def default_root_dir
@@ -544,7 +542,6 @@ module EasyML
 
     def new_model_file!
       build_model_file(
-        root_dir: root_dir,
         model: self,
         s3_bucket: EasyML::Configuration.s3_bucket,
         s3_region: EasyML::Configuration.s3_region,

data/app/models/easy_ml/model_file.rb

@@ -23,7 +23,7 @@ module EasyML
     belongs_to :model, class_name: "EasyML::Model"
 
     include EasyML::Concerns::Configurable
-    add_configuration_attributes :s3_bucket, :s3_prefix, :s3_region, :s3_access_key_id, :s3_secret_access_key, :root_dir
+    add_configuration_attributes :s3_bucket, :s3_prefix, :s3_region, :s3_access_key_id, :s3_secret_access_key
 
     def synced_file
       EasyML::Support::SyncedFile.new(
@@ -37,6 +37,21 @@ module EasyML
       )
     end
 
+    def root_dir
+      Pathname.new(model.root_dir)
+    end
+
+    def model_root
+      File.expand_path("..", root_dir.to_s)
+    end
+
+    def full_path(filename = nil)
+      filename = self.filename if filename.nil?
+      return nil if filename.nil?
+
+      root_dir.join(filename).to_s
+    end
+
     def exist?
       fit?
     end
@@ -54,33 +69,7 @@ module EasyML
 
     def upload(path)
       synced_file.upload(path)
-      set_path(path)
-    end
-
-    def set_path(path)
-      path = get_full_path(path)
-      basename = Pathname.new(path).basename.to_s
-      unless path.start_with?(full_dir)
-        new_path = File.join(full_dir, basename).to_s
-        FileUtils.mkdir_p(Pathname.new(new_path).dirname.to_s)
-        FileUtils.cp(path, new_path)
-        path = new_path
-      end
-      self.filename = basename
-      self.path = get_relative_path(path)
-    end
-
-    def get_full_path(path)
-      path = path.to_s
-
-      path = Rails.root.join(path) unless path.match?(Regexp.new(Rails.root.to_s))
-      path
-    end
-
-    def get_relative_path(path)
-      path = path.to_s
-      path = path.to_s.split(Rails.root.to_s).last
-      path.to_s.split("/")[0..-2].reject(&:empty?).join("/")
+      update(filename: Pathname.new(path).basename.to_s)
     end
 
     def download
@@ -94,26 +83,6 @@ module EasyML
       Digest::SHA256.file(full_path).hexdigest
     end
 
-    def full_path(filename = nil)
-      filename = self.filename if filename.nil?
-      return nil if filename.nil?
-      return nil if relative_dir.nil?
-
-      Rails.root.join(relative_dir, filename).to_s
-    end
-
-    def relative_dir
-      root_dir.to_s.gsub(Regexp.new(Rails.root.to_s), "").gsub!(%r{^/}, "")
-    end
-
-    def full_dir
-      Rails.root.join(relative_dir).to_s
-    end
-
-    def model_root
-      File.expand_path("..", full_dir)
-    end
-
     def cleanup!
       [model_root].each do |dir|
         EasyML::Support::FileRotate.new(dir, []).cleanup(extension_allowlist)

data/app/models/easy_ml/splitter_history.rb

@@ -1,3 +1,19 @@
+# == Schema Information
+#
+# Table name: easy_ml_splitter_histories
+#
+#  id                 :bigint           not null, primary key
+#  splitter_id        :integer          not null
+#  splitter_type      :string           not null
+#  configuration      :json
+#  dataset_id         :integer          not null
+#  created_at         :datetime         not null
+#  updated_at         :datetime         not null
+#  history_started_at :datetime         not null
+#  history_ended_at   :datetime
+#  history_user_id    :integer
+#  snapshot_id        :string
+#
 module EasyML
   class SplitterHistory < ActiveRecord::Base
     self.table_name = "easy_ml_splitter_histories"

data/app/serializers/easy_ml/prediction_serializer.rb

@@ -5,7 +5,12 @@ module EasyML
     include JSONAPI::Serializer
 
     attribute :prediction do |object|
-      object.prediction_value.symbolize_keys.dig(:value)
+      case object.prediction_value
+      when Hash
+        object.prediction_value.symbolize_keys.dig(:value)
+      when Numeric
+        object.prediction_value
+      end
     end
 
     attributes :id,

data/config/initializers/zhong.rb

@@ -7,5 +7,9 @@ if %w[zhong:start].include?(ARGV.first)
     every 1.hour, "cleanup" do
       EasyML::CleanJob.perform_later
     end
+
+    every 1.hour, "cleanup" do
+      EasyML::ScheduleRetrainingJob.perform_later
+    end
   end
 end

data/lib/easy_ml/data/date_converter.rb

@@ -3,6 +3,7 @@ module EasyML
     module DateConverter
       COMMON_DATE_FORMATS = [
         "%Y-%m-%dT%H:%M:%S.%6N",   # e.g., "2021-01-01T00:00:00.000000"
+        "%Y-%m-%d %H:%M:%S.%L Z", # e.g., "2025-01-03 23:04:49.492 Z"
         "%Y-%m-%d %H:%M:%S.%L",    # e.g., "2021-01-01 00:01:36.000"
         "%Y-%m-%d %H:%M:%S.%L",   # e.g., "2021-01-01 00:01:36.000"
         "%Y-%m-%d %H:%M:%S",      # e.g., "2021-01-01 00:01:36"

data/lib/easy_ml/data/polars_reader.rb

@@ -196,15 +196,22 @@ module EasyML
           polars_args[:dtypes].merge!(dtypes)
         end
         ext = Pathname.new(file).extname.gsub(/\./, "")
+        date_cols = []
         case ext
         when "csv"
-          filtered_args = filter_polars_args(Polars.method(:read_csv))
-          filtered_args.merge!(infer_schema_length: 1_000_000, null_values: ["\\N", "\\\\N", "NULL"])
+          filtered_args, date_cols = filter_polars_args(Polars.method(:read_csv))
+          filtered_args.merge!(
+            infer_schema_length: 1_000_000,
+            null_values: ["\\N", "\\\\N", "NULL"],
+          )
           df = Polars.read_csv(file, **filtered_args)
         when "parquet"
-          filtered_args = filter_polars_args(Polars.method(:read_parquet))
+          filtered_args, date_cols = filter_polars_args(Polars.method(:read_parquet))
           df = Polars.read_parquet(file, **filtered_args)
         end
+        date_cols.each do |col|
+          df = EasyML::Data::DateConverter.maybe_convert_date(df, col)
+        end
         df
       end
 
@@ -214,7 +221,13 @@ module EasyML
 
       def filter_polars_args(method)
         supported_params = method.parameters.map { |_, name| name }
-        polars_args.select { |k, _| supported_params.include?(k) }
+        filtered = polars_args.select { |k, _| supported_params.include?(k) }
+
+        # Filter out any datetime columns, and use maybe_convert_date to convert later
+        date_cols = (filtered[:dtypes] || {}).select { |k, v| v.class == Polars::Datetime }.keys
+        filtered[:dtypes] = (filtered[:dtypes] || {}).reject { |k, v| v.class == Polars::Datetime }.compact.to_h
+        filtered = filtered.select { |k, _| supported_params.include?(k) }
+        return filtered, date_cols
       end
 
       def csv_files

data/lib/easy_ml/data/statistics_learner.rb

@@ -59,7 +59,7 @@ module EasyML::Data
           stats[col].merge!(most_frequent_value: series.mode.sort.to_a&.first)
           if field_type == :categorical
             stats[col].merge!(
-              unique_count: series.n_unique,
+              unique_count: series.cast(:str).n_unique,
               counts: Hash[series.value_counts.to_hashes.map(&:values)],
             )
           end

data/lib/easy_ml/engine.rb

@@ -1,5 +1,6 @@
 require "aws-sdk"
 require "awesome_print"
+require "rails/all"
 require "inertia_rails"
 require "jsonapi/serializer"
 require "numo/narray"
@@ -68,6 +69,16 @@ module EasyML
       end
     end
 
+    initializer "easy_ml.check_pending_migrations" do
+      if defined?(Rails::Server)
+        config.after_initialize do
+          if EasyML.pending_migrations?
+            puts "\e[33mWARNING: You have pending EasyML migrations. Run 'rails generate easy_ml:migration' to add them.\e[0m"
+          end
+        end
+      end
+    end
+
     initializer "easy_ml.active_job_config" do
       resque_initializer = File.expand_path("config/initializers/resque.rb", root)
       require resque_initializer if File.exist?(resque_initializer)
@@ -77,6 +88,17 @@ 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.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/pending_migrations.rb

@@ -0,0 +1,19 @@
+module EasyML
+  def self.pending_migrations?
+    return false unless defined?(ActiveRecord)
+
+    # Get all migration files from our templates
+    template_dir = File.expand_path("../railtie/generators/templates/migration", __dir__)
+    template_migrations = Dir.glob(File.join(template_dir, "*.tt")).map do |f|
+      File.basename(f, ".tt").sub(/^create_/, "")
+    end
+
+    # Get all existing migrations
+    existing_migrations = Dir.glob(Rails.root.join("db/migrate/*_*.rb")).map do |f|
+      File.basename(f).sub(/^\d+_create_/, "").sub(/\.rb$/, "")
+    end
+
+    # Check if any template migrations are not in existing migrations
+    (template_migrations - existing_migrations).any?
+  end
+end

data/lib/easy_ml/predict.rb

@@ -10,25 +10,38 @@ module EasyML
       @models = {}
     end
 
-    def self.predict(model_name, df)
+    def self.predict(model_name, df, serialize: false)
       if df.is_a?(Hash)
         df = Polars::DataFrame.new(df)
       end
-      raw_input = df.to_hashes&.first
+      raw_input = df.to_hashes
       df = instance.normalize(model_name, df)
+      normalized_input = df.to_hashes
       preds = instance.predict(model_name, df)
       current_version = instance.get_model(model_name)
 
-      EasyML::Prediction.create!(
-        model: current_version.model,
-        model_history: current_version,
-        prediction_type: current_version.model.task,
-        prediction_value: {
-          value: preds.first,
-        }.compact,
-        raw_input: raw_input,
-        normalized_input: df.to_hashes&.first,
-      )
+      output = preds.zip(raw_input, normalized_input).map do |pred, raw, norm|
+        EasyML::Prediction.create!(
+          model: current_version.model,
+          model_history: current_version,
+          prediction_type: current_version.model.task,
+          prediction_value: pred,
+          raw_input: raw,
+          normalized_input: norm,
+        )
+      end
+
+      output = if output.is_a?(Array) && output.count == 1
+          output.first
+        else
+          output
+        end
+
+      if serialize
+        EasyML::PredictionSerializer.new(output).serializable_hash
+      else
+        output
+      end
     end
 
     def self.train(model_name, tuner: nil, evaluator: nil)