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

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e920c79d629c6ef5283aeb3f0e52e68618b881e1b83535e501663ca5fda0f693
-  data.tar.gz: 30b812805cd90b1e8159a31fb92fec46a6bcb61568bf6a542c04952a407301b7
+  metadata.gz: 6fc7f7afffd4e61312ab0302fb20e40a0eed3dcde05e16e9ea9ed1f523cb8681
+  data.tar.gz: d39511d86a3e01bc4ececcc20e1f7d0a583ac2eb07c8752c4ca2766b38318343
 SHA512:
-  metadata.gz: 9068e03f0ab50edaffb86d7403740d92257f9fd32ae97ca52b9dfda074cc8850e2a409d09c99fe1ed21735a3b6421b9187e4fa493c50943f1dc5d6ebc1bc7d7a
-  data.tar.gz: 9d326b950c8db91be5cd45745c13367d74fea713f68219afa8bab44153400d1f7b633e3962356e6c255604b20528f06b7d118775e3615d272d99453720dbe39d
+  metadata.gz: 81f8594d04fcaa7274e7fc1a20ff8f67bd03599988f43e8377d10734639eb3aa48b8516f0422c7ca496774c2922e1e6042105cea72a7bad01c3bf847a553f636
+  data.tar.gz: d5c7cf29044fe2ac184b88ae2b68f1b9c12550157c331f263eb2e3ee79225c40245736c2531610141138e58753ffeaa257145052c02aa6fbe7598cafb0699a78

data/Rakefile

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require "sprockets/railtie"
 require "bundler/gem_tasks"
 require "rspec/core/rake_task"
 
@@ -20,16 +19,19 @@ require_relative "lib/easy_ml"
 # Load the annotate tasks
 require "annotate/annotate_models"
 
+require "combustion"
+Combustion.path = "spec/internal"
+Combustion::Application.configure_for_combustion
 task :environment do
-  require "combustion"
-  require "sprockets"
-  Combustion.path = "spec/internal"
-  Combustion.initialize! :active_record do |config|
-    config.assets = ActiveSupport::OrderedOptions.new # Stub to avoid errors
-    config.assets.enabled = false # Set false since assets are handled by Vite
-  end
-  EasyML::Engine.eager_load!
+  Combustion::Application.initialize!
+
+  # Reset migrations paths so we can keep the migrations in the project root,
+  # not the Rails root
+  migrations_paths = ["spec/internal/db/migrate"]
+  ActiveRecord::Tasks::DatabaseTasks.migrations_paths = migrations_paths
+  ActiveRecord::Migrator.migrations_paths = migrations_paths
 end
+Combustion::Application.load_tasks
 
 namespace :easy_ml do
   task annotate_models: :environment do

data/app/controllers/easy_ml/settings_controller.rb

@@ -29,11 +29,10 @@ module EasyML
       EasyML::Configuration.configure do |config|
         config.storage = @settings.storage
         config.timezone = @settings.timezone
-        config.s3_access_key_id = @settings.s3_access_key_id
-        config.s3_secret_access_key = @settings.s3_secret_access_key
         config.s3_bucket = @settings.s3_bucket
         config.s3_region = @settings.s3_region
         config.s3_prefix = @settings.s3_prefix
+        config.wandb_api_key = @settings.wandb_api_key
       end
       flash.now[:notice] = "Settings saved."
       render inertia: "pages/SettingsPage", props: {
@@ -47,8 +46,6 @@ module EasyML
       params.require(:settings).permit(
         :storage,
         :timezone,
-        :s3_access_key_id,
-        :s3_secret_access_key,
         :s3_bucket,
         :s3_region,
         :s3_prefix,

data/app/frontend/pages/SettingsPage.tsx

@@ -1,7 +1,7 @@
 import React, { useState } from 'react';
 import { usePage } from '@inertiajs/react'
 import { useInertiaForm } from 'use-inertia-form';
-import { Settings2, Save, AlertCircle, Key, Database, Globe2 } from 'lucide-react';
+import { Settings2, Save, AlertCircle, Key, Globe2, Database } from 'lucide-react';
 import { PluginSettings } from '../components/settings/PluginSettings';
 
 interface Settings {
@@ -9,9 +9,6 @@ interface Settings {
     timezone: string;
     s3_bucket: string;
     s3_region: string;
-    s3_access_key_id: string;
-    s3_secret_access_key: string;
-    wandb_api_key: string;
   }
 }
 
@@ -88,7 +85,6 @@ export default function SettingsPage({ settings: initialSettings }: { settings:
               <select
                 id="timezone"
                 value={formData.settings.timezone}
-                
                 onChange={(e) => setFormData({
                   ...formData,
                   settings: {
@@ -113,7 +109,6 @@ export default function SettingsPage({ settings: initialSettings }: { settings:
           {/* S3 Configuration */}
           <div className="space-y-4">
             <div className="flex items-center gap-2 mb-4">
-
               <Database className="w-5 h-5 text-gray-500" />
               <h3 className="text-lg font-medium text-gray-900">S3 Configuration</h3>
             </div>
@@ -162,80 +157,6 @@ export default function SettingsPage({ settings: initialSettings }: { settings:
                 </select>
               </div>
             </div>
-
-            <div className="bg-blue-50 rounded-lg p-4">
-              <div className="flex gap-2">
-                <AlertCircle className="w-5 h-5 text-blue-500 mt-0.5" />
-                <div>
-                  <h4 className="text-sm font-medium text-blue-900">AWS Credentials</h4>
-                  <p className="mt-1 text-sm text-blue-700">
-                    These credentials will be used as default for all S3 operations. You can override them per datasource.
-                  </p>
-                </div>
-              </div>
-            </div>
-
-            <div className="grid grid-cols-2 gap-6">
-              <div>
-                <label htmlFor="accessKeyId" className="block text-sm font-medium text-gray-700 mb-1">
-                  Access Key ID
-                </label>
-                <div className="relative">
-                  <Key className="absolute left-3 top-1/2 transform -translate-y-1/2 w-4 h-4 text-gray-400" />
-                  <input
-                    type="text"
-                    id="accessKeyId"
-                    value={formData.settings.s3_access_key_id}
-                    onChange={(e) => setFormData({
-                      ...formData,
-                      settings: {
-                        ...formData.settings,
-                        s3_access_key_id: e.target.value
-                      }
-                    })}
-                    className="mt-1 block w-full pl-9 rounded-md border-gray-300 shadow-sm focus:border-blue-500 focus:ring-blue-500"
-                    placeholder="AKIA..."
-                  />
-                </div>
-              </div>
-
-              <div>
-                <label htmlFor="secretAccessKey" className="block text-sm font-medium text-gray-700 mb-1">
-                  Secret Access Key
-                </label>
-                <div className="relative">
-                  <Key className="absolute left-3 top-1/2 transform -translate-y-1/2 w-4 h-4 text-gray-400" />
-                  <input
-                    type={showSecretKey ? 'text' : 'password'}
-                    id="secretAccessKey"
-                    value={formData.settings.s3_secret_access_key}
-                    onChange={(e) => setFormData({
-                      ...formData,
-                      settings: {
-                        ...formData.settings,
-                        s3_secret_access_key: e.target.value
-                      }
-                    })}
-                    className="mt-1 block w-full pl-9 pr-24 rounded-md border-gray-300 shadow-sm focus:border-blue-500 focus:ring-blue-500"
-                    placeholder="Your secret key"
-                  />
-                  <button
-                    type="button"
-                    onClick={() => setShowSecretKey(!showSecretKey)}
-                    className="absolute right-2 top-1/2 transform -translate-y-1/2 text-sm text-gray-500 hover:text-gray-700"
-                  >
-                    {showSecretKey ? 'Hide' : 'Show'}
-                  </button>
-                </div>
-              </div>
-            </div>
-          </div>
-
-          <div className="border-t border-gray-200 pt-8">
-            <PluginSettings
-              settings={formData.settings}
-              setData={(settings) => setFormData({ ...settings })}
-            />
           </div>
 
           <div className="pt-6 border-t flex items-center justify-between">

data/app/jobs/easy_ml/batch_job.rb

@@ -12,7 +12,14 @@ module EasyML
       # E.g. EasyML::ComputeFeatureBatchJob.enqueue_batch(features.map(&:id))
       #
       def enqueue_batch(args_list, batch_id = default_batch_id)
-        args_list = args_list.map { |arg| arg.is_a?(Array) ? arg : [arg] }
+        args_list = args_list.map do |arg|
+          arg = arg.is_a?(Array) ? arg : [arg]
+          arg.map do |arg|
+            arg.merge!(
+              batch_id: batch_id,
+            )
+          end
+        end
         store_batch_arguments(batch_id, args_list)
 
         args_list.each do |args|
@@ -22,8 +29,45 @@ module EasyML
         batch_id
       end
 
+      def enqueue_ordered_batches(args_list)
+        parent_id = get_parent_batch_id(args_list)
+        store_batch_arguments(parent_id, args_list)
+
+        batch = args_list.first
+        rest = args_list[1..]
+
+        rest.map do |batch|
+          Resque.redis.rpush("batch:#{parent_id}:remaining", batch.to_json)
+        end
+
+        enqueue_batch(batch)
+      end
+
+      def enqueue_next_batch(caller, parent_id)
+        next_batch = Resque.redis.lpop("batch:#{parent_id}:remaining")
+        payload = Resque.decode(next_batch)
+
+        caller.enqueue_batch(payload)
+      end
+
+      def next_batch?(parent_id)
+        batches_remaining(parent_id) > 0
+      end
+
+      def batches_remaining(parent_id)
+        Resque.redis.llen("batch:#{parent_id}:remaining")
+      end
+
+      def cleanup_batch(parent_id)
+        Resque.redis.del("batch:#{parent_id}:remaining")
+      end
+
       private
 
+      def get_parent_batch_id(args_list)
+        args_list.dup.flatten.first.dig(:parent_batch_id)
+      end
+
       # Store batch arguments in Redis
       def store_batch_arguments(batch_id, args_list)
         redis_key = "#{batch(batch_id)}:original_args"

data/app/jobs/easy_ml/compute_feature_job.rb

@@ -1,19 +1,83 @@
 module EasyML
   class ComputeFeatureJob < BatchJob
+    extend EasyML::DataframeSerialization
+
     @queue = :easy_ml
 
     def self.perform(batch_id, options = {})
       options.symbolize_keys!
       feature_id = options.dig(:feature_id)
       feature = EasyML::Feature.find(feature_id)
-      feature.fit_batch(options)
+      dataset = feature.dataset
+
+      # Check if any feature has failed before proceeding
+      if dataset.features.any? { |f| f.workflow_status == "failed" }
+        puts "Aborting feature computation due to previous feature failure"
+        return
+      end
+
+      begin
+        feature.update(workflow_status: :analyzing) if feature.workflow_status == :ready
+        feature.fit_batch(options.merge!(batch_id: batch_id))
+      rescue => e
+        EasyML::Feature.transaction do
+          return if dataset.reload.workflow_status == :failed
+
+          feature.update(workflow_status: :failed)
+          dataset.update(workflow_status: :failed)
+          build_error_with_context(dataset, e, batch_id, feature)
+        end
+      end
+    end
+
+    def self.build_error_with_context(dataset, error, batch_id, feature)
+      error = EasyML::Event.handle_error(dataset, error)
+      batch = feature.build_batch(batch_id: batch_id)
+
+      # Convert any dataframes in the context to serialized form
+      error.create_context(context: batch)
     end
 
     def self.after_batch_hook(batch_id, *args)
       puts "After batch!"
-      feature_ids = fetch_batch_arguments(batch_id).flatten.map(&:symbolize_keys).pluck(:feature_id).uniq
-      dataset = EasyML::Feature.find_by(id: feature_ids.first).dataset
-      dataset.after_fit_features
+      batch_args = fetch_batch_arguments(batch_id).flatten.map(&:symbolize_keys)
+      feature_ids = batch_args.pluck(:feature_id).uniq
+      parent_id = batch_args.pluck(:parent_batch_id).first
+
+      feature = EasyML::Feature.find_by(id: feature_ids.first)
+
+      if feature.failed?
+        dataset.features.where(workflow_status: :analyzing).update_all(workflow_status: :ready)
+        return BatchJob.cleanup_batch(parent_id)
+      end
+
+      feature.after_fit
+
+      if BatchJob.next_batch?(parent_id)
+        BatchJob.enqueue_next_batch(self, parent_id)
+      else
+        dataset = EasyML::Feature.find_by(id: feature_ids.first).dataset
+        dataset.after_fit_features
+      end
+    end
+
+    private
+
+    def self.remove_remaining_batch_jobs(batch_id)
+      # Remove all remaining jobs in the batch
+      while (jobs = Resque.peek(:easy_ml, 0, 1000)).any?
+        jobs.each do |job|
+          if job["args"][0] == batch_id
+            Resque.dequeue(self, *job["args"])
+          end
+        end
+
+        # Break if we've processed all jobs (no more jobs match our batch_id)
+        break unless jobs.any? { |job| job["args"][0] == batch_id }
+      end
     end
   end
 end
+
+# If any feature fails, the entire batch fails
+# If any feature fails, the RELATED batches should fail

data/app/models/concerns/easy_ml/dataframe_serialization.rb

@@ -0,0 +1,30 @@
+module EasyML
+  module DataframeSerialization
+    extend ActiveSupport::Concern
+
+    def serialize_dataframe(df)
+      return unless df
+      JSON.parse(df.write_json)
+    end
+
+    def deserialize_dataframe(df_data)
+      return unless df_data.present? && df_data.key?("columns")
+
+      columns = df_data["columns"].map do |col|
+        dtype = case col["datatype"]
+          when Hash
+            if col["datatype"]["Datetime"]
+              Polars::Datetime.new(col["datatype"]["Datetime"][0].downcase.to_sym).class
+            else
+              Polars::Utf8
+            end
+          else
+            Polars.const_get(col["datatype"])
+          end
+        Polars::Series.new(col["name"], col["values"], dtype: dtype)
+      end
+
+      Polars::DataFrame.new(columns)
+    end
+  end
+end

data/app/models/easy_ml/dataset.rb

@@ -102,13 +102,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 destructively_cleanup!
@@ -219,8 +215,11 @@ module EasyML
     end
 
     def after_fit_features
-      features.update_all(needs_fit: false, fit_at: Time.current)
       unlock!
+      reload
+      return if failed?
+
+      features.update_all(needs_fit: false, fit_at: Time.current)
       actually_refresh
     end
 
@@ -281,22 +280,24 @@ module EasyML
     end
 
     def refreshing
-      return false if is_history_class?
-      unlock! unless analyzing?
-
-      lock_dataset do
-        update(workflow_status: "analyzing")
-        fully_reload
-        yield
-      ensure
-        unlock!
-      end
-    rescue => e
-      update(workflow_status: "failed")
-      e.backtrace.grep(/easy_ml/).each do |line|
-        puts line
+      begin
+        return false if is_history_class?
+        unlock! unless analyzing?
+
+        lock_dataset do
+          update(workflow_status: "analyzing")
+          fully_reload
+          yield
+        ensure
+          unlock!
+        end
+      rescue => e
+        update(workflow_status: "failed")
+        e.backtrace.grep(/easy_ml/).each do |line|
+          puts line
+        end
+        raise e
       end
-      raise e
     end
 
     def unlock!

data/app/models/easy_ml/dataset_history.rb

@@ -30,17 +30,12 @@ module EasyML
     self.table_name = "easy_ml_dataset_histories"
     include Historiographer::History
 
-    has_many :columns,
-      ->(dataset_history) { where(snapshot_id: dataset_history.snapshot_id) },
+    has_many :columns, ->(dataset_history) { where(snapshot_id: dataset_history.snapshot_id) },
       class_name: "EasyML::ColumnHistory",
       foreign_key: "dataset_id",
       primary_key: "dataset_id",
       extend: EasyML::ColumnList
 
-    def root_dir
-      read_attribute(:root_dir)
-    end
-
     def fit
       false
     end

data/app/models/easy_ml/datasources/polars_datasource.rb

@@ -1,6 +1,8 @@
 module EasyML
   module Datasources
     class PolarsDatasource < BaseDatasource
+      include EasyML::DataframeSerialization
+
       validates :df, presence: true
       add_configuration_attributes :df
 
@@ -58,7 +60,7 @@ module EasyML
         return unless df
 
         datasource.configuration = (datasource.configuration || {}).merge(
-          "df" => JSON.parse(df.write_json),
+          "df" => serialize_dataframe(df),
         )
       end
 
@@ -66,23 +68,7 @@ module EasyML
         return unless datasource.configuration&.key?("df")
 
         df_data = datasource.configuration["df"]
-        return unless df_data.present? && df_data.key?("columns")
-
-        columns = df_data["columns"].map do |col|
-          dtype = case col["datatype"]
-            when Hash
-              if col["datatype"]["Datetime"]
-                Polars::Datetime.new(col["datatype"]["Datetime"][0].downcase.to_sym).class
-              else
-                Polars::Utf8
-              end
-            else
-              Polars.const_get(col["datatype"])
-            end
-          Polars::Series.new(col["name"], col["values"], dtype: dtype)
-        end
-
-        datasource.df = Polars::DataFrame.new(columns)
+        datasource.df = deserialize_dataframe(df_data)
       end
     end
   end

data/app/models/easy_ml/event.rb

@@ -19,6 +19,7 @@ module EasyML
     STATUSES = %w[started success failed].freeze
 
     belongs_to :eventable, polymorphic: true, optional: true
+    has_one :context, dependent: :destroy, class_name: "EasyML::EventContext"
 
     validates :name, presence: true
     validates :status, presence: true, inclusion: { in: STATUSES }
@@ -51,8 +52,8 @@ module EasyML
           error = e
         end
       end
-      create_event(model, "failed", error)
       Rails.logger.error("#{self.class.name} failed: #{error.message}")
+      create_event(model, "failed", error)
     end
 
     def self.format_stacktrace(error)

data/app/models/easy_ml/event_context.rb

@@ -0,0 +1,58 @@
+# == Schema Information
+#
+# Table name: easy_ml_event_contexts
+#
+#  id         :bigint           not null, primary key
+#  event_id   :bigint           not null
+#  context    :jsonb            not null
+#  created_at :datetime         not null
+#  updated_at :datetime         not null
+#
+module EasyML
+  class EventContext < ActiveRecord::Base
+    include EasyML::DataframeSerialization
+
+    self.table_name = "easy_ml_event_contexts"
+
+    belongs_to :event
+
+    validates :context, presence: true
+    validates :event, presence: true
+
+    def context=(new_context)
+      write_attribute(:context, serialize_context(new_context))
+      @context = new_context
+    end
+
+    def context
+      @context ||= deserialize_context(read_attribute(:context))
+    end
+
+    private
+
+    def serialize_context(new_context)
+      case new_context
+      when Hash
+        self.format = :json
+        new_context.to_json
+      when YAML
+        self.format = :yaml
+        new_context.to_yaml
+      when Polars::DataFrame
+        self.format = :dataframe
+        serialize_dataframe(new_context)
+      end
+    end
+
+    def deserialize_context(context)
+      case format.to_sym
+      when :json
+        JSON.parse(context)
+      when :yaml
+        YAML.safe_load(context)
+      when :dataframe
+        deserialize_dataframe(context)
+      end
+    end
+  end
+end