Always use environment variables from add time + minor changes (#140)

* store environment variables at add time; made a bunch of things faster; better logging

* do not call scancel if the job list is empty

* optimize hash query, use multiprocessing for uploading src files

* remove redundant typing

* proper config updates in src reload

* use multithreading instead of multiprocessing

* cache reading file configs

* make sure that args are hashable

* fix setting experiments

* do not cache files from the venv

Author: n-gao

.pre-commit-config.yaml

@@ -60,6 +60,7 @@ repos:
           - setuptools>=69.2.0
           - importlib_resources>=5.7.0
           - typing_extensions>=4.10
+          - deepdiff>=7.0.1
           - ruff>=0.6.1
           - pytest>=8.3.2
           - pre-commit>=3.8.0

examples/example_config.yaml

@@ -60,18 +60,25 @@ seml:
   description: An example configuration.
 
 slurm:
+  - experiments_per_job: 1
+    sbatch_options:
+      gres: gpu:0
+      mem: 1G
+      cpus-per-task: 1
+      time: 0-08:00
+      partition: cpu_all,cpu_large
   - experiments_per_job: 4
     sbatch_options:
       gres: gpu:1 # num GPUs
-      mem: 16G # memory
-      cpus-per-task: 2 # num cores
+      mem: 1G # memory
+      cpus-per-task: 1 # num cores
       time: 0-08:00 # max time, D-HH:MM
       partition: gpu_gtx1080
   - experiments_per_job: 16
     sbatch_options:
       gres: gpu:1 # num GPUs
-      mem: 16G # memory
-      cpus-per-task: 2 # num cores
+      mem: 1G # memory
+      cpus-per-task: 1 # num cores
       time: 0-08:00 # max time, D-HH:MM
       partition: gpu_a100
 

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     "setuptools>=69.2.0",
     "importlib_resources>=5.7.0",
     "typing_extensions>=4.10",
+    "deepdiff>=7.0.1",
 ]
 
 [project.optional-dependencies]

src/seml/commands/add.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import copy
-import datetime
 import logging
 import os
 from typing import TYPE_CHECKING, Any, Dict, cast
@@ -16,6 +15,7 @@
     read_config,
     remove_duplicates,
     remove_prepended_dashes,
+    requires_interpolation,
     resolve_configs,
     resolve_interpolations,
 )
@@ -29,6 +29,7 @@
     remove_keys_from_nested,
     to_typeddict,
     unflatten,
+    utcnow,
 )
 from seml.utils.errors import ConfigError
 
@@ -125,7 +126,7 @@ def add_configs(
                 **{
                     '_id': start_id + idx,
                     'status': States.STAGED[0],
-                    'add_time': datetime.datetime.utcnow(),
+                    'add_time': utcnow(),
                 },
             },
         )
@@ -134,11 +135,13 @@ def add_configs(
     if description is not None:
         for db_dict in documents:
             db_dict['seml']['description'] = description
-    if resolve_descriptions:
-        for db_dict in documents:
-            if 'description' in db_dict['seml']:
+        # If description is not supplied via CLI, it will already be resolved.
+        if resolve_descriptions and requires_interpolation(
+            {'description': description}, ['description']
+        ):
+            for db_dict in documents:
                 db_dict['seml']['description'] = resolve_description(
-                    db_dict['seml']['description'], db_dict
+                    db_dict['seml'].get('description', ''), db_dict
                 )
 
     collection.insert_many(documents)
@@ -324,7 +327,6 @@ def add_config_file(
                     }
                 ),
                 'config_unresolved': config_unresolved,
-                'seml': seml_config,
             },
         )
         for config, config_unresolved in zip(configs, configs_unresolved)
@@ -342,7 +344,7 @@ def add_config_file(
     for document in documents:
         document['config_hash'] = make_hash(
             document['config'],
-            config_get_exclude_keys(document['config'], document['config_unresolved']),
+            config_get_exclude_keys(document['config_unresolved']),
         )
 
     if not force_duplicates:

src/seml/commands/manage.py

@@ -402,13 +402,16 @@ def delete_experiments(
     """
     from seml.console import prompt
 
+    collection = get_collection(db_collection_name)
     # Before deleting, we should first cancel the experiments that are still running.
     if cancel:
         cancel_states = set(States.PENDING + States.RUNNING)
         if filter_states is not None and len(filter_states) > 0:
             cancel_states = cancel_states.intersection(filter_states)
 
-        if len(cancel_states) > 0:
+        if len(cancel_states) > 0 and collection.find_one(
+            build_filter_dict(cancel_states, batch_id, filter_dict, sacred_id)
+        ):
             cancel_experiments(
                 db_collection_name,
                 sacred_id,
@@ -420,12 +423,9 @@ def delete_experiments(
                 wait=True,
             )
 
-    collection = get_collection(db_collection_name)
     experiment_files_to_delete = []
 
-    filter_dict = build_filter_dict(
-        filter_states, batch_id, filter_dict, sacred_id=sacred_id
-    )
+    filter_dict = build_filter_dict(filter_states, batch_id, filter_dict, sacred_id)
     ndelete = collection.count_documents(filter_dict)
     if sacred_id is not None and ndelete == 0:
         raise MongoDBError(f'No experiment found with ID {sacred_id}.')
@@ -560,7 +560,12 @@ def get_experiment_reset_op(exp: ExperimentDoc):
     ]
 
     # Clean up SEML dictionary
-    keep_seml = {'source_files', 'working_dir', SETTINGS.SEML_CONFIG_VALUE_VERSION}
+    keep_seml = {
+        'source_files',
+        'working_dir',
+        'env',
+        SETTINGS.SEML_CONFIG_VALUE_VERSION,
+    }
     keep_seml.update(SETTINGS.VALID_SEML_CONFIG_VALUES)
     seml_keys = set(exp['seml'].keys())
     for key in seml_keys - keep_seml:
@@ -627,6 +632,9 @@ def reset_experiments(
     exps = collection.find(filter_dict)
     if sacred_id is not None and nreset == 0:
         raise MongoDBError(f'No experiment found with ID {sacred_id}.')
+    if nreset == 0:
+        logging.info('No experiments to reset.')
+        return
 
     logging.info(f'Resetting the state of {nreset} experiment{s_if(nreset)}.')
     if nreset >= SETTINGS.CONFIRM_THRESHOLD.RESET:
@@ -761,7 +769,8 @@ def reload_sources(
     """
     from importlib.metadata import version
 
-    import gridfs
+    from bson import ObjectId
+    from deepdiff import DeepDiff
     from pymongo import UpdateOne
 
     from seml.console import prompt
@@ -774,7 +783,15 @@ def reload_sources(
         filter_dict = {}
     db_results = list(
         collection.find(
-            filter_dict, {'batch_id', 'seml', 'config', 'status', 'config_unresolved'}
+            filter_dict,
+            {
+                'batch_id',
+                'seml',
+                'config',
+                'status',
+                'config_unresolved',
+                'config_hash',
+            },
         )
     )
     id_to_document: dict[int, list[ExperimentDoc]] = {}
@@ -845,29 +862,40 @@ def reload_sources(
             )
         ]
 
-        result = collection.bulk_write(
-            [
-                UpdateOne(
-                    {'_id': document['_id']},
-                    {
-                        '$set': {
-                            'config': document['config'],
-                            'config_unresolved': document['config_unresolved'],
-                            'config_hash': make_hash(
-                                document['config'],
-                                config_get_exclude_keys(
-                                    document['config'], document['config_unresolved']
-                                ),
-                            ),
-                        }
-                    },
+        # determine which documents to udpate
+        updates = []
+        for old_doc, new_doc in zip(documents, new_documents):
+            use_hash = 'config_hash' in old_doc
+            # these config fields are populated if the experiment ran
+            runtime_fields = {
+                k: old_doc['config'][k]
+                for k in ['db_collection', 'overwrite', 'seed']
+                if k in old_doc['config']
+            }
+            new = dict(
+                config=new_doc['config'] | runtime_fields,
+                config_unresolved=new_doc['config_unresolved'],
+            )
+            # compare new to old config
+            if use_hash:
+                new['config_hash'] = make_hash(
+                    new_doc['config'],
+                    config_get_exclude_keys(new_doc['config_unresolved']),
                 )
-                for document in new_documents
-            ]
-        )
-        logging.info(
-            f'Batch {batch_id}: Resolved configurations of {result.matched_count} experiments against new source files ({result.modified_count} changed).'
-        )
+                update = new['config_hash'] != old_doc['config_hash']
+            else:
+                diff = DeepDiff(new['config'], old_doc['config'], ignore_order=True)
+                update = bool(diff)
+            # Create mongodb update
+            if update:
+                updates.append(UpdateOne({'_id': old_doc['_id']}, {'$set': new}))
+        if len(updates) > 0:
+            result = collection.bulk_write(updates)
+            logging.info(
+                f'Batch {batch_id}: Resolved configurations of {result.matched_count} experiments against new source files ({result.modified_count} changed).'
+            )
+        else:
+            logging.info(f'Batch {batch_id}: No experiment configurations changed.')
 
         # Check whether the configurations aligns with the current source code
         check_config(
@@ -879,7 +907,6 @@ def reload_sources(
 
         # Find the currently used source files
         db = collection.database
-        fs = gridfs.GridFS(db)
         fs_filter_dict = {
             'metadata.batch_id': batch_id,
             'metadata.collection_name': f'{collection.name}',
@@ -915,8 +942,7 @@ def reload_sources(
         except Exception as e:
             logging.error(f'Batch {batch_id}: Failed to set new source files.')
             # Delete new source files from DB
-            for to_delete in source_files:
-                fs.delete(to_delete[1])
+            delete_files(db, [x[1] for x in source_files])
             raise e
 
         # Delete the old source files
@@ -927,10 +953,10 @@ def reload_sources(
                 'metadata.deprecated': True,
             }
             source_files_old = [
-                x['_id'] for x in db['fs.files'].find(fs_filter_dict, {'_id'})
+                cast(ObjectId, x['_id'])
+                for x in db['fs.files'].find(fs_filter_dict, {'_id'})
             ]
-            for to_delete in source_files_old:
-                fs.delete(to_delete)
+            delete_files(db, source_files_old)
 
 
 def detect_duplicates(