Repository: mlgym/mlgym (language: python)

Issue:
Accelerate multi gpu support

149 server sends mlgym event messages in json format to client
Waiting for PR #147 , before marking this PR as ready, and afterwards create a new PR for #152 


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index c18ebcce..40fecb18 100644
--- a/.gitignore
+++ b/.gitignore
@@ -126,3 +126,6 @@ src/ml_gym/frontend/node_modules
 src/ml_gym/backend/websocket_api/event_storage
 event_storage
 *.whl
+*.arrow
+transformer/
+
diff --git a/README.md b/README.md
index ebea265c..0a2e14a8 100644
--- a/README.md
+++ b/README.md
@@ -76,7 +76,6 @@ The parameter `process_count` specifies the number of experiments that we run in
 cd mlgym/example/grid_search_example
 
 python run.py --process_count 3 \
-              --text_logging_path general_logging/ \
               --num_epochs 10 \
               --websocket_logging_servers http://127.0.0.1:5002 \
               --gs_rest_api_endpoint http://127.0.0.1:5001 \
diff --git a/example/bert/bert_lm_blueprint.py b/example/bert/bert_lm_blueprint.py
new file mode 100644
index 00000000..f9407118
--- /dev/null
+++ b/example/bert/bert_lm_blueprint.py
@@ -0,0 +1,100 @@
+from ml_gym.error_handling.exception import DatasetNotFoundError
+from typing import Dict, List, Any
+from ml_gym.data_handling.postprocessors.factory import ModelGymInformedIteratorFactory
+from ml_gym.modes import RunMode
+import torch
+from ml_gym.blueprints.constructables import ComponentConstructable, ModelRegistryConstructable
+from ml_gym.blueprints.blue_prints import BluePrint
+from ml_gym.batching.batch import DatasetBatch
+from dataclasses import dataclass
+from ml_gym.blueprints.component_factory import ComponentFactory, Injector
+from ml_gym.data_handling.postprocessors.collator import Collator
+from bert_model import BERTLM
+from transformers import DataCollatorForLanguageModeling, BertTokenizerFast
+from data_stack.dataset.meta import MetaFactory
+from data_stack.dataset.iterator import InformedDatasetIteratorIF
+from datasets import load_from_disk
+from mlm_loss_function import LMLossFunctionRegistryConstructable
+
+
+@dataclass
+class LMWikiBookCorpusDatasetConstructable(ComponentConstructable):
+    dataset_identifier: str = ""
+    dataset_folder_path: str = None
+
+    def _construct_impl(self) -> Dict[str, InformedDatasetIteratorIF]:
+        if self.dataset_folder_path is None:
+            raise DatasetNotFoundError("Dataset path not specified")
+        chunked_tokenized_dataset = load_from_disk(self.dataset_folder_path)
+        split_name = "train"
+        iterator_meta = MetaFactory.get_iterator_meta(sample_pos=0, target_pos=1, tag_pos=0)
+        dataset_meta = MetaFactory.get_dataset_meta(identifier=self.component_identifier,
+                                                    dataset_name=self.dataset_identifier,
+                                                    dataset_tag=split_name,
+                                                    iterator_meta=iterator_meta)
+        dataset_dict = {split_name: ModelGymInformedIteratorFactory.get_dataset_iterator(chunked_tokenized_dataset, dataset_meta)}
+        return dataset_dict
+
+
+@dataclass
+class BERRTLMCollator(Collator):
+
+    def __init__(self, target_publication_key: str, tokenizer_file_path: str, pad_to_multiple_of: int = 8, mlm_probability: float = 0.15):
+        self.target_publication_key = target_publication_key
+        tokenizer = BertTokenizerFast(tokenizer_file=tokenizer_file_path)  # "trained_wiki_tokenizer/tokenizer.json"
+        self.data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer,
+                                                             mlm_probability=mlm_probability,
+                                                             pad_to_multiple_of=pad_to_multiple_of)
+
+    def __call__(self, batch: List[torch.Tensor]) -> DatasetBatch:
+        """
+        :param batch: batch format [no_samples, height, width, channels]
+        :return:
+        """
+        collated_batch = self.data_collator(batch)
+        samples = collated_batch["input_ids"]
+        targets = {self.target_publication_key: collated_batch["labels"],
+                   "attention_mask": collated_batch["attention_mask"]}
+        return DatasetBatch(targets=targets, tags=None, samples=samples)
+
+
+@dataclass
+class MyModelRegistryConstructable(ModelRegistryConstructable):
+    def _construct_impl(self):
+        super()._construct_impl()
+        self.model_registry.add_class("lm_bert", BERTLM)
+        return self.model_registry
+
+
+class BERTLMBluePrint(BluePrint):
+    def __init__(self, run_mode: RunMode, config: Dict, grid_search_id: str,
+                 experiment_id: str, external_injection: Dict[str, Any] = None,
+                 warm_start_epoch: int = 0):
+        super().__init__(run_mode, config, grid_search_id, experiment_id, external_injection, warm_start_epoch)
+
+    @staticmethod
+    def construct_components(config: Dict, component_names: List[str], device: torch.device,
+                             external_injection: Dict[str, Any] = None) -> Dict[str, Any]:
+        if external_injection is not None:
+            injection_mapping = {"id_bert_lm_collator": BERRTLMCollator,
+                                 "id_computation_device": device,
+                                 **external_injection}
+            injector = Injector(injection_mapping, raise_mapping_not_found=True)
+
+        else:
+            injection_mapping = {"id_bert_lm_collator": BERRTLMCollator}
+            injector = Injector(injection_mapping, raise_mapping_not_found=False)
+
+        component_factory = ComponentFactory(injector)
+        component_factory.register_component_type("MODEL_REGISTRY", "DEFAULT", MyModelRegistryConstructable)
+        component_factory.register_component_type("DATASET_ITERATORS", "LMWikiBookCorpusDataset", LMWikiBookCorpusDatasetConstructable),
+        component_factory.register_component_type("LOSS_FUNCTION_REGISTRY", "LM", LMLossFunctionRegistryConstructable)
+
+        components = component_factory.build_components_from_config(config, component_names)
+        return components
+
+    def construct(self, device: torch.device = None) -> Dict[str, Any]:
+        component_names = ["model", "trainer", "optimizer", "evaluator", "early_stopping_strategy", "checkpointing_strategy", "lr_scheduler"]
+        components = BERTLMBluePrint.construct_components(self.config, component_names, device, self.external_injection)
+
+        return components
diff --git a/example/bert/bert_lm_config.yml b/example/bert/bert_lm_config.yml
new file mode 100644
index 00000000..5de5bf9c
--- /dev/null
+++ b/example/bert/bert_lm_config.yml
@@ -0,0 +1,223 @@
+global_config:
+  storage_connector_path: &storage_path_anchor ./file_storage/
+  seed: &seed_value 2
+  target_key: &target_key_anchor target_key
+  model_prediction_key: &model_prediction_key_anchor model_prediction_key
+  postprocessing_argmax_key: &postprocessing_argmax_key_anchor postprocessing_argmax_key
+
+dataset_iterators:
+  component_type_key: DATASET_ITERATORS
+  variant_key: LMWikiBookCorpusDataset
+  config:
+    dataset_identifier: LMWikiBookCorpus
+    dataset_folder_path: /scratch/max/mlgym/example/bert/chunked_tokenized_dataset_train
+
+split_dataset_iterators:
+  component_type_key: SPLITTED_DATASET_ITERATORS
+  variant_key: RANDOM
+  requirements:
+    - name: iterators
+      component_name: dataset_iterators
+      subscription:
+        - train
+  config:
+    split_configs:
+      train:
+        val: 0.001
+        test: 0.001
+        train_lm: 0.998
+    seed: 2
+
+split_dataset_iterator_2:
+  component_type_key: SPLITTED_DATASET_ITERATORS
+  variant_key: RANDOM
+  requirements:
+    - name: iterators
+      component_name: split_dataset_iterators
+      subscription: [train_lm, val, test]
+  config:
+    split_configs:
+      train_lm:
+        rest: 0.999
+        train_small: 0.001
+    seed: 3
+
+data_collator:
+  component_type_key: DATA_COLLATOR
+  variant_key: DEFAULT
+  config:
+    collator_type:
+      injectable:
+        id: id_bert_lm_collator
+    collator_params:
+      target_publication_key: *target_key_anchor
+      tokenizer_file_path: /scratch/max/mlgym/example/bert/tokenizers/trained_wiki_tokenizer/tokenizer.json
+      pad_to_multiple_of: 8
+      mlm_probability: 0.15
+
+data_loaders:
+  component_type_key: DATA_LOADER
+  variant_key: DEFAULT
+  requirements:
+    - name: iterators
+      component_name: split_dataset_iterator_2
+      subscription: [train_lm, train_small, val, test]
+    - name: data_collator
+      component_name: data_collator
+  config:
+    batch_size: 20
+    seeds:
+      train_lm: 1
+      train_small: 2
+      val: 3
+      test: 4
+
+model_registry:
+  component_type_key: MODEL_REGISTRY
+  variant_key: DEFAULT
+
+model:
+  component_type_key: MODEL
+  variant_key: DEFAULT
+  requirements:
+    - name: model_registry
+      component_name: model_registry
+      subscription: lm_bert
+  config:
+    prediction_publication_keys:
+      prediction_publication_key: *model_prediction_key_anchor
+
+optimizer:
+  component_type_key: OPTIMIZER
+  variant_key: DEFAULT
+  config:
+    optimizer_key: ADAMW
+    params:
+      lr:
+        sweep: absolute
+        values: [5.e-5]
+
+lr_scheduler:
+  component_type_key: LR_SCHEDULER
+  variant_key: DEFAULT
+  config:
+    lr_scheduler_key: hugging_face_scheduler
+    params:
+      name: linear
+      num_warmup_steps: 50
+      num_training_steps: 1000
+
+loss_function_registry:
+  component_type_key: LOSS_FUNCTION_REGISTRY
+  variant_key: LM
+
+metric_registry:
+  component_type_key: METRIC_REGISTRY
+  variant_key: DEFAULT
+
+prediction_postprocessing_registry:
+  component_type_key: PREDICTION_POSTPROCESSING_REGISTRY
+  variant_key: DEFAULT
+
+train_component:
+  component_type_key: TRAIN_COMPONENT
+  variant_key: ACCELERATE
+  requirements:
+    - name: loss_function_registry
+      component_name: loss_function_registry
+    - name: prediction_postprocessing_registry
+      component_name: prediction_postprocessing_registry
+  config:
+    loss_fun_config:
+      prediction_subscription_key: *model_prediction_key_anchor
+      target_subscription_key: *target_key_anchor
+      key: MLMCrossEntropyLoss
+      tag: mlm_entropy_loss
+      vocab_size: 30522
+
+trainer:
+  component_type_key: TRAINER
+  variant_key: ACCELERATE
+  requirements:
+    - name: train_component
+      component_name: train_component
+    - name: model
+      component_name: model
+      subscription: null
+    - name: data_loaders
+      component_name: data_loaders
+      subscription: train_lm
+  config: {}
+
+eval_component:
+  component_type_key: EVAL_COMPONENT
+  variant_key: ACCELERATE
+  requirements:
+    - name: model
+      component_name: model
+      subscription: null
+    - name: data_loaders
+      component_name: data_loaders
+      subscription: [train_small, val, test]
+    - name: loss_function_registry
+      component_name: loss_function_registry
+    - name: metric_registry
+      component_name: metric_registry
+    - name: prediction_postprocessing_registry
+      component_name: prediction_postprocessing_registry
+  config:
+    cpu_target_subscription_keys:
+      - *target_key_anchor
+    cpu_prediction_subscription_keys: []
+      # - *postprocessing_argmax_key_anchor
+      # - *model_prediction_key_anchor
+    # post_processors_config:
+    #   - key: "ARG_MAX"
+    #     prediction_subscription_key: *model_prediction_key_anchor
+    #     prediction_publication_key: *postprocessing_argmax_key_anchor
+    metrics_config: []
+    loss_funs_config:
+      - prediction_subscription_key: *model_prediction_key_anchor
+        target_subscription_key: *target_key_anchor
+        key: MLMCrossEntropyLoss
+        tag: mlm_entropy_loss
+        vocab_size: 30522
+
+evaluator:
+  component_type_key: EVALUATOR
+  variant_key: ACCELERATE
+  requirements:
+    - name: eval_component
+      component_name: eval_component
+
+early_stopping_strategy_registry:
+  component_type_key: EARLY_STOPPING_STRATEGY_REGISTRY
+  variant_key: DEFAULT
+
+early_stopping_strategy:
+  component_type_key: EARLY_STOPPING_STRATEGY
+  variant_key: DEFAULT
+  requirements:
+    - name: early_stopping_strategy_registry
+      component_name: early_stopping_strategy_registry
+  config:
+    early_stopping_key: LAST_K_EPOCHS_IMPROVEMENT_STRATEGY
+    early_stopping_config:
+      min_relative_improvement: 0.00001
+      epochs_window: 50
+      split_name: val
+      monitoring_key: "mlm_entropy_loss"
+      is_increase_task: false
+
+checkpointing_strategy_registry:
+  component_type_key: CHECKPOINTING_STRATEGY_REGISTRY
+  variant_key: DEFAULT
+
+checkpointing_strategy:
+  component_type_key: CHECKPOINTING_STRATEGY
+  variant_key: DEFAULT
+  requirements:
+    - name: checkpointing_strategy_registry
+      component_name: checkpointing_strategy_registry
+  config:
+    checkpointing_key: SAVE_LAST_EPOCH_ONLY_CHECKPOINTING_STRATEGY
diff --git a/example/bert/bert_model.py b/example/bert/bert_model.py
new file mode 100644
index 00000000..dd99538d
--- /dev/null
+++ b/example/bert/bert_model.py
@@ -0,0 +1,41 @@
+from ml_gym.models.nn.net import NNModel
+from typing import Dict
+import torch
+from transformers import AutoConfig, AutoModelForMaskedLM
+
+
+class BERTLM(NNModel):
+
+    def __init__(self, prediction_publication_key: str, bert_version: str = "bert-base-uncased"):
+        super().__init__()
+        self.prediction_publication_key = prediction_publication_key
+        config = AutoConfig.from_pretrained(bert_version)
+        self.model = AutoModelForMaskedLM.from_config(config)
+
+    def forward_impl(self, inputs: torch.Tensor) -> Dict[str, torch.Tensor]:
+        outputs = self.model(inputs)
+        output_dict = {self.prediction_publication_key: outputs.logits}
+        return output_dict
+
+    def forward(self, inputs: torch.Tensor) -> Dict[str, torch.Tensor]:
+        return self.forward_impl(inputs)
+
+
+if __name__ == '__main__':
+    from datasets import load_from_disk
+    from transformers import DataCollatorForLanguageModeling, BertTokenizerFast
+    from torch.utils.data import DataLoader
+
+    tokenizer = BertTokenizerFast(tokenizer_file="/scratch/max/mlgym/example/transformer/tokenizers/trained_wiki_tokenizer/tokenizer.json")
+    chunked_tokenized_dataset = load_from_disk("example/transformer/preprocessed_datasets/chunked_tokenized_dataset_train")
+    mlm_probability = 0.15
+    batch_size = 30
+
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=mlm_probability, pad_to_multiple_of=8)
+
+    train_dataloader = DataLoader(chunked_tokenized_dataset, shuffle=True, batch_size=batch_size, collate_fn=data_collator)
+
+    sample = next(iter(train_dataloader))
+    model = BERTLM("")
+    prediction = model(sample["input_ids"])
+    print("")
diff --git a/example/bert/mlm_loss_function.py b/example/bert/mlm_loss_function.py
new file mode 100644
index 00000000..310fb34e
--- /dev/null
+++ b/example/bert/mlm_loss_function.py
@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import Dict, Any
+from ml_gym.loss_functions.loss_functions import Loss
+from ml_gym.batching.batch import InferenceResultBatch
+import torch
+from torch.nn import CrossEntropyLoss
+from ml_gym.blueprints.constructables import LossFunctionRegistryConstructable
+
+
+@dataclass
+class LMLossFunctionRegistryConstructable(LossFunctionRegistryConstructable):
+    class LossKeys:
+        MLMCrossEntropyLoss = "MLMCrossEntropyLoss"
+
+    def _construct_impl(self):
+        loss_fun_registry = super()._construct_impl()
+        default_mapping: Dict[str, Any] = {
+            self.LossKeys.MLMCrossEntropyLoss: MLMCrossEntropyLoss
+        }
+
+        for key, loss_type in default_mapping.items():
+            loss_fun_registry.add_class(key, loss_type)
+
+        return loss_fun_registry
+
+
+class MLMCrossEntropyLoss(Loss):
+    def __init__(self, target_subscription_key: str, prediction_subscription_key: str, vocab_size: int,
+                 tag: str = ""):
+        super().__init__(tag)
+        self.target_subscription_key = target_subscription_key
+        self.prediction_subscription_key = prediction_subscription_key
+        self.vocab_size = vocab_size
+        self.loss_fun = CrossEntropyLoss()
+
+    def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor:
+        t = forward_batch.get_targets(self.target_subscription_key)
+        p = forward_batch.get_predictions(self.prediction_subscription_key)
+        masked_lm_loss = self.loss_fun(p.view(-1, self.vocab_size), t.view(-1))
+        return masked_lm_loss
diff --git a/example/bert/run.py b/example/bert/run.py
new file mode 100644
index 00000000..cb25ffb4
--- /dev/null
+++ b/example/bert/run.py
@@ -0,0 +1,10 @@
+
+from bert_lm_blueprint import BERTLMBluePrint
+from ml_gym.cmd_entrypoint.cmd import get_args, run
+
+
+if __name__ == '__main__':
+
+    blueprint_class = BERTLMBluePrint
+    config_path = get_args()
+    run(blueprint_class=blueprint_class, run_configuration_file_path=config_path)
diff --git a/example/bert/run_config.yml b/example/bert/run_config.yml
new file mode 100644
index 00000000..38cde87f
--- /dev/null
+++ b/example/bert/run_config.yml
@@ -0,0 +1,14 @@
+run_configuration:
+  type: train # train, warmstart
+  config:
+    num_epochs: 1000 # Number of epochs
+    num_batches_per_epoch: 100
+    gs_config_path: /scratch/max/mlgym/example/bert/bert_lm_config.yml
+
+environment:
+  type: accelerate # multiprocessing,  main_process, accelerate
+
+logging:
+  websocket_logging_servers: # List of websocket logging servers, e.g., http://127.0.0.1:9090 http://127.0.0.1:8080
+    - http://127.0.0.1:5002
+  gs_rest_api_endpoint: http://127.0.0.1:5001 # Endpoint for the grid search API, e.g., http://127.0.0.1:8080
diff --git a/example/bert/train_bert.sh b/example/bert/train_bert.sh
new file mode 100644
index 00000000..dda5948e
--- /dev/null
+++ b/example/bert/train_bert.sh
@@ -0,0 +1,11 @@
+ 
+#!/bin/sh
+
+
+python run.py --process_count 1 \
+              --gpus 7 \
+              --num_epochs 3 \
+              --websocket_logging_servers http://127.0.0.1:5002 \
+              --gs_rest_api_endpoint http://127.0.0.1:5001 \
+              train \
+              --gs_config_path /scratch/max/mlgym/example/bert/bert_lm_config.yml
diff --git a/example/bert/train_bert_accelerate.sh b/example/bert/train_bert_accelerate.sh
new file mode 100644
index 00000000..5577e14e
--- /dev/null
+++ b/example/bert/train_bert_accelerate.sh
@@ -0,0 +1,5 @@
+ 
+#!/bin/sh
+
+
+accelerate launch --multi_gpu --num_processes=5 --gpu_ids 0,1,5,6,7 run.py --config_path /scratch/max/mlgym/example/bert/run_config.yml
diff --git a/example/cnn_ddp/cnn_experiment_config.yml b/example/cnn_ddp/cnn_experiment_config.yml
new file mode 100644
index 00000000..787359f5
--- /dev/null
+++ b/example/cnn_ddp/cnn_experiment_config.yml
@@ -0,0 +1,251 @@
+global_config:
+  storage_connector_path: &storage_path_anchor ./file_storage/
+  seed: &seed_value 2
+  target_key: &target_key_anchor target_key
+  model_prediction_key: &model_prediction_key_anchor model_prediction_key
+  postprocessing_argmax_key: &postprocessing_argmax_key_anchor postprocessing_argmax_key
+
+dataset_repository:
+  component_type_key: DATASET_REPOSITORY
+  variant_key: DEFAULT
+  config:
+    storage_connector_path: *storage_path_anchor
+
+dataset_iterators:
+  component_type_key: DATASET_ITERATORS
+  variant_key: DEFAULT
+  requirements:
+    - name: repository
+      component_name: dataset_repository
+  config:
+    dataset_identifier: mnist
+    split_configs:
+      - split: train
+      - split: test
+
+splitted_dataset_iterators:
+  component_type_key: SPLITTED_DATASET_ITERATORS
+  variant_key: RANDOM
+  requirements:
+    - name: iterators
+      component_name: dataset_iterators
+      subscription:
+        - train
+        - test
+  config:
+    split_configs:
+      train:
+        train: 0.7
+        val: 0.3
+    seed: 2
+
+data_collator:
+  component_type_key: DATA_COLLATOR
+  variant_key: DEFAULT
+  config:
+    collator_type:
+      injectable:
+        id: id_conv_mnist_standard_collator
+    collator_params:
+      target_publication_key: *target_key_anchor
+
+data_loaders:
+  component_type_key: DATA_LOADER
+  variant_key: DEFAULT
+  requirements:
+    - name: iterators
+      component_name: splitted_dataset_iterators
+      subscription: [train, val, test]
+    - name: data_collator
+      component_name: data_collator
+  config:
+    batch_size: 50
+    weigthed_sampling_split_name: train
+    seeds:
+      train: 0
+      val: 1
+      test: 2
+
+model_registry:
+  component_type_key: MODEL_REGISTRY
+  variant_key: DEFAULT
+
+model:
+  component_type_key: MODEL
+  variant_key: DEFAULT
+  requirements:
+    - name: model_registry
+      component_name: model_registry
+      subscription: conv_net
+  config:
+    model_definition:
+      layer_config:
+        - params:
+            in_channels: 1
+            kernel_size: 3
+            out_channels: 32
+            stride: 1
+          type: conv
+        - params:
+            in_channels: 32
+            kernel_size: 3
+            out_channels: 64
+            stride: 1
+          type: conv
+        - params:
+            in_features: 9216
+            out_features: 128
+          type: fc
+        - params:
+            in_features: 128
+            out_features: 10
+          type: fc
+    prediction_publication_keys:
+      prediction_publication_key: *model_prediction_key_anchor
+    seed: *seed_value
+
+optimizer:
+  component_type_key: OPTIMIZER
+  variant_key: DEFAULT
+  config:
+    optimizer_key: ADAM
+    params:
+      lr: 
+        sweep: absolute
+        values: [1, 0.1, 0.01, 0.001, 0.0001, 0.00001] # [1, 0.1, 0.01, 0.001, 0.0001]
+
+lr_scheduler:
+  component_type_key: LR_SCHEDULER
+  variant_key: DEFAULT
+  config:
+    lr_scheduler_key: dummy
+
+loss_function_registry:
+  component_type_key: LOSS_FUNCTION_REGISTRY
+  variant_key: DEFAULT
+
+metric_registry:
+  component_type_key: METRIC_REGISTRY
+  variant_key: DEFAULT
+
+prediction_postprocessing_registry:
+  component_type_key: PREDICTION_POSTPROCESSING_REGISTRY
+  variant_key: DEFAULT
+
+train_component:
+  component_type_key: TRAIN_COMPONENT
+  variant_key: ACCELERATE
+  requirements:
+    - name: loss_function_registry
+      component_name: loss_function_registry
+    - name: prediction_postprocessing_registry
+      component_name: prediction_postprocessing_registry
+  config:
+    loss_fun_config:
+      prediction_subscription_key: *model_prediction_key_anchor
+      target_subscription_key: *target_key_anchor
+      key: CrossEntropyLoss
+      tag: cross_entropy_loss
+
+trainer:
+  component_type_key: TRAINER
+  variant_key: ACCELERATE
+  requirements:
+    - name: train_component
+      component_name: train_component
+    - name: model
+      component_name: model
+      subscription: null
+    - name: data_loaders
+      component_name: data_loaders
+      subscription: train
+
+eval_component:
+  component_type_key: EVAL_COMPONENT
+  variant_key: ACCELERATE
+  requirements:
+    - name: model
+      component_name: model
+      subscription: null
+    - name: data_loaders
+      component_name: data_loaders
+      subscription: [train, val, test]
+    - name: loss_function_registry
+      component_name: loss_function_registry
+    - name: metric_registry
+      component_name: metric_registry
+    - name: prediction_postprocessing_registry
+      component_name: prediction_postprocessing_registry
+  config:
+    cpu_target_subscription_keys:
+      - *target_key_anchor
+    cpu_prediction_subscription_keys:
+      - *postprocessing_argmax_key_anchor
+      - *model_prediction_key_anchor
+    post_processors_config:
+      - key: "ARG_MAX"
+        prediction_subscription_key: *model_prediction_key_anchor
+        prediction_publication_key: *postprocessing_argmax_key_anchor
+    metrics_config:
+      - key: F1_SCORE
+        params:
+          average: macro
+        prediction_subscription_key: *postprocessing_argmax_key_anchor
+        target_subscription_key: *target_key_anchor
+        tag: F1_SCORE_macro
+      - key: PRECISION
+        params:
+          average: macro
+        prediction_subscription_key: *postprocessing_argmax_key_anchor
+        target_subscription_key: *target_key_anchor
+        tag: PRECISION_macro
+      - key: RECALL
+        params:
+          average: macro
+        prediction_subscription_key: *postprocessing_argmax_key_anchor
+        target_subscription_key: *target_key_anchor
+        tag: RECALL_macro
+    loss_funs_config:
+      - prediction_subscription_key: *model_prediction_key_anchor
+        target_subscription_key: *target_key_anchor
+        key: CrossEntropyLoss
+        tag: cross_entropy_loss
+
+evaluator:
+  component_type_key: EVALUATOR
+  variant_key: ACCELERATE
+  requirements:
+    - name: eval_component
+      component_name: eval_component
+
+early_stopping_strategy_registry:
+  component_type_key: EARLY_STOPPING_STRATEGY_REGISTRY
+  variant_key: DEFAULT
+
+early_stopping_strategy:
+  component_type_key: EARLY_STOPPING_STRATEGY
+  variant_key: DEFAULT
+  requirements:
+    - name: early_stopping_strategy_registry
+      component_name: early_stopping_strategy_registry
+  config:
+    early_stopping_key: LAST_K_EPOCHS_IMPROVEMENT_STRATEGY
+    early_stopping_config:
+      min_relative_improvement: 0.00001
+      epochs_window: 5
+      split_name: val
+      monitoring_key: "F1_SCORE_macro"
+      is_increase_task: true
+
+checkpointing_strategy_registry:
+  component_type_key: CHECKPOINTING_STRATEGY_REGISTRY
+  variant_key: DEFAULT
+
+checkpointing_strategy:
+  component_type_key: CHECKPOINTING_STRATEGY
+  variant_key: DEFAULT
+  requirements:
+    - name: checkpointing_strategy_registry
+      component_name: checkpointing_strategy_registry
+  config:
+    checkpointing_key: SAVE_LAST_EPOCH_ONLY_CHECKPOINTING_STRATEGY
diff --git a/example/cnn_ddp/conv_net.py b/example/cnn_ddp/conv_net.py
new file mode 100644
index 00000000..189b176f
--- /dev/null
+++ b/example/cnn_ddp/conv_net.py
@@ -0,0 +1,52 @@
+from ml_gym.models.nn.net import NNModel
+from typing import Dict
+from torch import nn
+import torch
+import torch.nn.functional as F
+
+
+class ConvNet(NNModel):
+
+    def __init__(self, prediction_publication_key: str, layer_config: Dict, seed: int = 0):
+        super().__init__(seed=seed)
+        self.prediction_publication_key = prediction_publication_key
+        self.conv_layers = nn.ModuleList([])
+        self.fc_layers = nn.ModuleList([])
+        for layer in layer_config:
+            if layer["type"] == "conv":
+                self.conv_layers.append(ConvNet.create_conv_layer_from_config(layer["params"]))
+            elif layer["type"] == "fc":
+                self.fc_layers.append(ConvNet.create_fc_layer_from_config(layer["params"]))
+
+    @staticmethod
+    def create_conv_layer_from_config(layer_dict) -> nn.Module:
+        return nn.Conv2d(in_channels=layer_dict["in_channels"],
+                         out_channels=layer_dict["out_channels"],
+                         kernel_size=layer_dict["kernel_size"],
+                         stride=layer_dict["stride"])
+
+    @staticmethod
+    def create_fc_layer_from_config(layer_dict) -> nn.Module:
+        return nn.Linear(in_features=layer_dict["in_features"],
+                         out_features=layer_dict["out_features"])
+
+    def forward_impl(self, inputs: torch.Tensor) -> Dict[str, torch.Tensor]:
+        output = inputs
+        output = self.conv_layers[0](output)
+        output = F.relu(output)
+
+        output = self.conv_layers[1](output)
+        output = F.relu(output)
+        output = F.max_pool2d(output, 2, 2)
+        output = F.dropout(output, p=0.25, training=True)
+
+        output = output.view(inputs.shape[0], -1)
+        output = self.fc_layers[0](output)
+        output = F.relu(output)
+        output = F.dropout(output, p=0.5, training=True)
+        output = self.fc_layers[1](output)
+
+        return {self.prediction_publication_key: output}
+
+    def forward(self, inputs: torch.Tensor) -> Dict[str, torch.Tensor]:
+        return self.forward_impl(inputs)
diff --git a/example/cnn_ddp/conv_net_blueprint.py b/example/cnn_ddp/conv_net_blueprint.py
new file mode 100644
index 00000000..ee740104
--- /dev/null
+++ b/example/cnn_ddp/conv_net_blueprint.py
@@ -0,0 +1,71 @@
+from typing import Dict, List, Any
+from ml_gym.modes import RunMode
+import torch
+from ml_gym.blueprints.constructables import ModelRegistryConstructable
+from ml_gym.blueprints.blue_prints import BluePrint
+from ml_gym.batching.batch import DatasetBatch
+from dataclasses import dataclass
+from ml_gym.blueprints.component_factory import ComponentFactory, Injector
+from ml_gym.data_handling.postprocessors.collator import Collator
+from conv_net import ConvNet
+
+
+@dataclass
+class MNISTCollator(Collator):
+    target_publication_key: str = None
+
+    def __call__(self, batch: List[torch.Tensor]):
+        """
+        :param batch: batch format [no_samples, height, width, channels]
+        :return:
+        """
+        # batch contains a list of tuples of structure (sequence, target)
+        inputs = [item[0].to(self.device) for item in batch]
+        inputs = torch.stack(inputs)
+        # transform into vector
+        # inputs = inputs.view(inputs.shape[0], -1)
+        # transform into CHW matrix
+        # inputs = inputs.permute(0, 3, 1, 2)
+        inputs = inputs.reshape(-1, 1, 28, 28)
+        targets_tensor = torch.tensor([item[1] for item in batch]).to(inputs[0].device)
+        target_partitions = {self.target_publication_key: targets_tensor}
+        return DatasetBatch(targets=target_partitions, tags=None, samples=inputs)
+
+
+@dataclass
+class MyModelRegistryConstructable(ModelRegistryConstructable):
+    def _construct_impl(self):
+        super()._construct_impl()
+        self.model_registry.add_class("conv_net", ConvNet)
+        return self.model_registry
+
+
+class ConvNetBluePrint(BluePrint):
+    def __init__(self, run_mode: RunMode, config: Dict, grid_search_id: str,
+                 experiment_id: str, external_injection: Dict[str, Any] = None, warm_start_epoch: int = 0):
+        super().__init__(run_mode, config, grid_search_id, experiment_id, external_injection, warm_start_epoch)
+
+    @staticmethod
+    def construct_components(config: Dict, component_names: List[str], device: torch.device,
+                             external_injection: Dict[str, Any] = None) -> Dict[str, Any]:
+        if external_injection is not None:
+            injection_mapping = {"id_conv_mnist_standard_collator": MNISTCollator,
+                                 "id_computation_device": device,
+                                 **external_injection}
+            injector = Injector(injection_mapping, raise_mapping_not_found=True)
+
+        else:
+            injection_mapping = {"id_conv_mnist_standard_collator": MNISTCollator}
+            injector = Injector(injection_mapping, raise_mapping_not_found=False)
+
+        component_factory = ComponentFactory(injector)
+        component_factory.register_component_type("MODEL_REGISTRY", "DEFAULT", MyModelRegistryConstructable)
+
+        components = component_factory.build_components_from_config(config, component_names)
+        return components
+
+    def construct(self, device: torch.device = None) -> Dict[str, Any]:
+        component_names = ["model", "trainer", "optimizer", "lr_scheduler", "evaluator", "early_stopping_strategy", "checkpointing_strategy"]
+        components = ConvNetBluePrint.construct_components(self.config, component_names, device, self.external_injection)
+
+        return components
diff --git a/example/cnn_ddp/run copy.py b/example/cnn_ddp/run copy.py
new file mode 100644
index 00000000..a3a65c34
--- /dev/null
+++ b/example/cnn_ddp/run copy.py
@@ -0,0 +1,8 @@
+from ml_gym.cmd_entrypoint.cmd import get_args, run
+from conv_net_blueprint import ConvNetBluePrint
+
+if __name__ == '__main__':
+
+    blueprint_class = ConvNetBluePrint
+    config_path = get_args()
+    run(blueprint_class=blueprint_class, run_configuration_file_path=config_path)
diff --git a/example/cnn_ddp/run.py b/example/cnn_ddp/run.py
new file mode 100644
index 00000000..d46b6df3
--- /dev/null
+++ b/example/cnn_ddp/run.py
@@ -0,0 +1,10 @@
+
+from conv_net_blueprint import ConvNetBluePrint
+from ml_gym.cmd_entrypoint.cmd import get_args, run
+
+
+if __name__ == '__main__':
+
+    blueprint_class = ConvNetBluePrint
+    config_path = get_args()
+    run(blueprint_class=blueprint_class, run_configuration_file_path=config_path)
diff --git a/example/cnn_ddp/run_config.yml b/example/cnn_ddp/run_config.yml
new file mode 100644
index 00000000..2c6ee53e
--- /dev/null
+++ b/example/cnn_ddp/run_config.yml
@@ -0,0 +1,14 @@
+run_configuration:
+  type: train # train, warmstart
+  config:
+    num_epochs: 10 # Number of epochs
+    # num_batches_per_epoch: 100
+    gs_config_path: /scratch/max/mlgym/example/cnn_ddp/cnn_experiment_config.yml
+
+environment:
+  type: accelerate # multiprocessing,  main_process, accelerate
+
+logging:
+  websocket_logging_servers: # List of websocket logging servers, e.g., http://127.0.0.1:9090 http://127.0.0.1:8080
+    - http://127.0.0.1:5002
+  gs_rest_api_endpoint: http://127.0.0.1:5001 # Endpoint for the grid search API, e.g., http://127.0.0.1:8080
diff --git a/example/cnn_ddp/train_cnn_accelerate.sh b/example/cnn_ddp/train_cnn_accelerate.sh
new file mode 100644
index 00000000..6f0ec452
--- /dev/null
+++ b/example/cnn_ddp/train_cnn_accelerate.sh
@@ -0,0 +1,5 @@
+ 
+#!/bin/sh
+
+
+accelerate launch --multi_gpu --num_processes=8 --gpu_ids 0,1,2,3,4,5,6,7 run.py --config_path /scratch/max/mlgym/example/cnn_ddp/run_config.yml
diff --git a/example/cnn_ddp/train_cnn_accelerate_small.sh b/example/cnn_ddp/train_cnn_accelerate_small.sh
new file mode 100644
index 00000000..04cc52b0
--- /dev/null
+++ b/example/cnn_ddp/train_cnn_accelerate_small.sh
@@ -0,0 +1,5 @@
+ 
+#!/bin/sh
+
+
+accelerate launch --multi_gpu --num_processes=2 --gpu_ids 4,5 run.py --config_path /scratch/max/mlgym/example/cnn_ddp/run_config.yml
diff --git a/example/grid_search_example/conv_net_blueprint.py b/example/grid_search_example/conv_net_blueprint.py
index 968f4a2b..ee740104 100644
--- a/example/grid_search_example/conv_net_blueprint.py
+++ b/example/grid_search_example/conv_net_blueprint.py
@@ -1,15 +1,12 @@
 from typing import Dict, List, Any
 from ml_gym.modes import RunMode
-from ml_gym.persistency.logging import ExperimentStatusLogger, MLgymStatusLoggerCollectionConstructable
 import torch
 from ml_gym.blueprints.constructables import ModelRegistryConstructable
 from ml_gym.blueprints.blue_prints import BluePrint
-from ml_gym.gym.jobs import AbstractGymJob, GymJobFactory
 from ml_gym.batching.batch import DatasetBatch
 from dataclasses import dataclass
 from ml_gym.blueprints.component_factory import ComponentFactory, Injector
 from ml_gym.data_handling.postprocessors.collator import Collator
-from ml_gym.persistency.io import GridSearchAPIClientConstructableIF
 from conv_net import ConvNet
 
 
@@ -44,16 +41,13 @@ def _construct_impl(self):
 
 
 class ConvNetBluePrint(BluePrint):
-    def __init__(self, run_mode: RunMode, config: Dict, num_epochs: int, grid_search_id: str,
-                 gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-                 experiment_id: str, external_injection: Dict[str, Any] = None,
-                 logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None,
-                 warm_start_epoch: int = 0):
-        super().__init__(run_mode, num_epochs, config, grid_search_id, gs_api_client_constructable,
-                         experiment_id, external_injection, logger_collection_constructable, warm_start_epoch)
+    def __init__(self, run_mode: RunMode, config: Dict, grid_search_id: str,
+                 experiment_id: str, external_injection: Dict[str, Any] = None, warm_start_epoch: int = 0):
+        super().__init__(run_mode, config, grid_search_id, experiment_id, external_injection, warm_start_epoch)
 
     @staticmethod
-    def construct_components(config: Dict, component_names: List[str], device: torch.device, external_injection: Dict[str, Any] = None) -> Dict[str, Any]:
+    def construct_components(config: Dict, component_names: List[str], device: torch.device,
+                             external_injection: Dict[str, Any] = None) -> Dict[str, Any]:
         if external_injection is not None:
             injection_mapping = {"id_conv_mnist_standard_collator": MNISTCollator,
                                  "id_computation_device": device,
@@ -70,20 +64,8 @@ def construct_components(config: Dict, component_names: List[str], device: torch
         components = component_factory.build_components_from_config(config, component_names)
         return components
 
-    def construct(self, device: torch.device = None) -> AbstractGymJob:
+    def construct(self, device: torch.device = None) -> Dict[str, Any]:
         component_names = ["model", "trainer", "optimizer", "lr_scheduler", "evaluator", "early_stopping_strategy", "checkpointing_strategy"]
         components = ConvNetBluePrint.construct_components(self.config, component_names, device, self.external_injection)
 
-        logger_collection = self.logger_collection_constructable.construct()
-        experiment_status_logger = ExperimentStatusLogger(logger=logger_collection, grid_search_id=self.grid_search_id,
-                                                          experiment_id=self.experiment_id)
-
-        gym_job = GymJobFactory.get_gym_job(run_mode=self.run_mode,
-                                            grid_search_id=self.grid_search_id,
-                                            experiment_id=self.experiment_id,
-                                            num_epochs=self.num_epochs,
-                                            warm_start_epoch=self.warm_start_epoch,
-                                            experiment_status_logger=experiment_status_logger,
-                                            gs_api_client=self.gs_api_client_constructable.construct(),
-                                            **components)
-        return gym_job
+        return components
diff --git a/example/grid_search_example/gs_config.yml b/example/grid_search_example/gs_config.yml
index 30285164..cafd492b 100644
--- a/example/grid_search_example/gs_config.yml
+++ b/example/grid_search_example/gs_config.yml
@@ -186,7 +186,6 @@ eval_component:
       - key: "ARG_MAX"
         prediction_subscription_key: *model_prediction_key_anchor
         prediction_publication_key: *postprocessing_argmax_key_anchor
-    train_split_name: train
     metrics_config:
       - key: F1_SCORE
         params:
@@ -211,7 +210,7 @@ eval_component:
         target_subscription_key: *target_key_anchor
         key: CrossEntropyLoss
         tag: cross_entropy_loss
-    show_progress: true
+    show_progress: false
 
 evaluator:
   component_type_key: EVALUATOR
diff --git a/example/grid_search_example/run.py b/example/grid_search_example/run.py
index 6ce8ccc6..a3a65c34 100644
--- a/example/grid_search_example/run.py
+++ b/example/grid_search_example/run.py
@@ -1,150 +1,8 @@
-import argparse
-from ml_gym.persistency.logging import MLgymStatusLoggerCollectionConstructable, LoggerConstructableIF, MLgymStatusLoggerConfig, \
-    MLgymStatusLoggerTypes
-from ml_gym.modes import RunMode, ValidationMode
+from ml_gym.cmd_entrypoint.cmd import get_args, run
 from conv_net_blueprint import ConvNetBluePrint
-from ml_gym.starter import mlgym_entry_train, mlgym_entry_warm_start
-from ml_gym.validation.validator_factory import get_validator
-from ml_gym.io.config_parser import YAMLConfigLoader
-from typing import List
-from ml_gym.persistency.io import GridSearchAPIClientConfig, GridSearchAPIClientConstructableIF, GridSearchAPIClientConstructable, \
-    GridSearchAPIClientType
-from pathlib import Path
-
-
-def get_logger_constructable(websocket_logging_servers: List) -> LoggerConstructableIF:
-    """
-    Create config object for websocket server used for logging events.
-
-    :params: websocket_logging_servers (list) : socket server config from CLI ``Ex: --websocket_logging_servers``
-
-    """
-    websocket_logging_servers = [(f"{protocol}:{ip}", int(port)) for protocol, ip, port in [
-        connection.split(":") for connection in websocket_logging_servers]]
-
-    logging_configs = [MLgymStatusLoggerConfig(logger_type=MLgymStatusLoggerTypes.STREAMED_LOGGER, logger_config={"host": ip, "port": port})
-                       for ip, port in websocket_logging_servers]
-    logger_collection_constructable = MLgymStatusLoggerCollectionConstructable(logging_configs)
-    return logger_collection_constructable
-
-
-def get_grid_search_restful_api_client_constructable(endpoint: str) -> GridSearchAPIClientConstructableIF:
-    """
-    Create config object for grid search RestAPI used for communicating with websocket server using HTTP APIs.
-
-    :params: endpoint (str) : rest API endpoint config from CLI ``Ex: --gs_rest_api_endpoint``
-
-    """
-    client_config = GridSearchAPIClientConfig(api_client_type=GridSearchAPIClientType.GRID_SEARCH_RESTFUL_API_CLIENT,
-                                              api_client_config={"endpoint": endpoint})
-    client_constructable = GridSearchAPIClientConstructable(client_config)
-    return client_constructable
-
-
-def entry_train(args):
-    """
-    Start mlGym and create job which starts grid search from scratch.
-
-    :params: args : config from CLI ``Ex: --process_count``
-
-    """
-    logger_collection_constructable = get_logger_constructable(args.websocket_logging_servers)
-    gs_restful_api_client_constructable = get_grid_search_restful_api_client_constructable(endpoint=args.gs_rest_api_endpoint)
-    blueprint_class = ConvNetBluePrint
-
-    gs_config_string = Path(args.gs_config_path).read_text()
-    gs_config = YAMLConfigLoader.load_string(gs_config_string)
-
-    if args.validation_strategy_config_path is not None:
-        validation_strategy_config_string = Path(args.validation_strategy_config_path).read_text()
-        validation_strategy_config = YAMLConfigLoader.load(validation_strategy_config_string)
-        validation_mode = ValidationMode[list(validation_strategy_config.keys())[0]]
-    else:
-        validation_strategy_config_string = None
-        validation_strategy_config = None
-        validation_mode = ValidationMode.GRID_SEARCH
-
-    validator = get_validator(validation_mode, blueprint_class, RunMode.TRAIN, validation_strategy_config, gs_config)
-
-    mlgym_entry_train(blueprint_class=blueprint_class,
-                      logger_collection_constructable=logger_collection_constructable,
-                      gs_api_client_constructable=gs_restful_api_client_constructable,
-                      gs_config_raw_string=gs_config_string,
-                      validation_strategy_config_raw_string=validation_strategy_config_string,
-                      validator=validator,
-                      text_logging_path=args.text_logging_path,
-                      process_count=args.process_count,
-                      gpus=args.gpus,
-                      log_std_to_file=args.log_std_to_file,
-                      num_epochs=args.num_epochs)
-
-
-def entry_warm_start(args):
-    """
-    Start mlGym and create job which resumes from previously started grid search.
-
-    :params: args : config from CLI ``Ex: --process_count``
-
-    """
-    blueprint_class = ConvNetBluePrint
-    logger_collection_constructable = get_logger_constructable(args.websocket_logging_servers)
-    gs_restful_api_client_constructable = get_grid_search_restful_api_client_constructable(endpoint=args.gs_rest_api_endpoint)
-
-    mlgym_entry_warm_start(blueprint_class=blueprint_class,
-                           grid_search_id=args.grid_search_id,
-                           logger_collection_constructable=logger_collection_constructable,
-                           gs_api_client_constructable=gs_restful_api_client_constructable,
-                           text_logging_path=args.text_logging_path,
-                           process_count=args.process_count,
-                           gpus=args.gpus,
-                           log_std_to_file=args.log_std_to_file,
-                           num_epochs=args.num_epochs)
-
-    # starter = MLGymTrainWarmStarter(grid_search_id=grid_search_id, blueprint_class=blueprint_class, validation_mode=)
-
-
-def parse_args_and_run():
-    """
-    Extract args from Command line and decides whether to train model from scratch or resume from previously started grid search.
-
-    Extracts websocket and grid search API from command line and create configuration to communicate with it for logging in event_storage path.
-
-    """
-    parser = argparse.ArgumentParser(description='Run a grid search on CPUs or distributed over multiple GPUs')
-
-    subparsers = parser.add_subparsers(help='Mutually exclusive arguments for TRAIN and WARM_START')
-
-    # Train
-    parser_train = subparsers.add_parser('train', help='Trains the models from scratch')
-    parser_train.set_defaults(func=entry_train)
-    parser_train.add_argument('--gs_config_path', type=str, required=True, help='Path to the grid search config')
-    parser_train.add_argument('--validation_strategy_config_path', type=str, required=False, help='Path to the validation strategy config')
-    parser_train.add_argument('--early_stopping_config_path', type=str, required=False, help='Path to the early stopping config')
-
-    # Warmstart
-    parser_warm_start = subparsers.add_parser('warm_start', help='Starts off from a previously started grid search')
-    parser_warm_start.set_defaults(func=entry_warm_start)
-    parser_warm_start.add_argument('--grid_search_id', type=str, required=True, help='Gridsearch id identifying the specific grid search')
-    parser_warm_start.add_argument('--early_stopping_config_path', type=str, required=False, help='Path to the early stopping config')
-
-    # parser.add_argument('--run_mode', choices=['TRAIN', 'WARM_START'], required=True)
-
-    # Common
-    parser.add_argument('--num_epochs', type=int, required=True, help='Number of epoch')
-    parser.add_argument('--process_count', type=int, required=True, help='Max. number of processes running at a time.')
-    parser.add_argument('--text_logging_path', type=str, required=True, help='Path to python textual logging directory')
-    parser.add_argument('--gpus', type=int, nargs='+', help='Indices of GPUs to distribute the GS over', default=None)
-    parser.add_argument('--log_std_to_file', default=False, action="store_true", help='Flag for forwarding std output to file')
-    parser.add_argument('--websocket_logging_servers', type=str, nargs='+',
-                        help='List of websocket logging servers, e.g., 127.0.0.1:9090 127.0.0.1:8080', default=None)
-    parser.add_argument('--gs_rest_api_endpoint', type=str, required=True,
-                        help='Endpoint for the grid search API, e.g., 127.0.0.1:8080')
-    # parser.add_argument('--validation_mode', choices=['NESTED_CROSS_VALIDATION', 'GRID_SEARCH', 'CROSS_VALIDATION'], required=True)
-
-    args = parser.parse_args()
-    args.func(args)
-
 
 if __name__ == '__main__':
 
-    parse_args_and_run()
+    blueprint_class = ConvNetBluePrint
+    config_path = get_args()
+    run(blueprint_class=blueprint_class, run_configuration_file_path=config_path)
diff --git a/example/grid_search_example/run_config.yml b/example/grid_search_example/run_config.yml
new file mode 100644
index 00000000..22c1fd62
--- /dev/null
+++ b/example/grid_search_example/run_config.yml
@@ -0,0 +1,17 @@
+run_configuration:
+  type: train # train, warmstart
+  config:
+    num_epochs: 10 # Number of epochs
+    gs_config_path: /scratch/max/mlgym/example/grid_search_example/gs_config.yml
+
+environment:
+  type: multiprocessing # multiprocessing,  main_process, accelerate
+  config:
+    process_count: 1 # Max. number of processes running at a time.
+    computation_device_ids: [4] # Indices of GPUs to distribute the GS over
+
+logging:
+  websocket_logging_servers: # List of websocket logging servers, e.g., http://127.0.0.1:9090 http://127.0.0.1:8080
+    - http://127.0.0.1:5002
+  gs_rest_api_endpoint: http://127.0.0.1:5001 # Endpoint for the grid search API, e.g., http://127.0.0.1:8080
+
diff --git a/example/grid_search_example/train_gs_mnist_conv_net_websocket.sh b/example/grid_search_example/train_gs_mnist_conv_net_websocket.sh
index f05a6280..008d1ba9 100644
--- a/example/grid_search_example/train_gs_mnist_conv_net_websocket.sh
+++ b/example/grid_search_example/train_gs_mnist_conv_net_websocket.sh
@@ -2,9 +2,8 @@
 #!/bin/sh
 
 
-python run.py --process_count 1 \
-              --text_logging_path general_logging/ \
-              --num_epochs 10 \
+python run.py --process_count 21 \
+              --num_epochs 100 \
               --websocket_logging_servers http://127.0.0.1:5002 \
               --gs_rest_api_endpoint http://127.0.0.1:5001 \
               train \
diff --git a/example/legacy/conv_net_blueprint.py b/example/legacy/conv_net_blueprint.py
index 4813e445..984e8c9c 100644
--- a/example/legacy/conv_net_blueprint.py
+++ b/example/legacy/conv_net_blueprint.py
@@ -4,7 +4,7 @@
 import torch
 from ml_gym.blueprints.constructables import ModelRegistryConstructable
 from ml_gym.blueprints.blue_prints import BluePrint
-from ml_gym.gym.jobs import AbstractGymJob, GymJobFactory
+from ml_gym.gym.gym_jobs.standard_gym_job import AbstractGymJob, GymJobFactory
 from ml_gym.batching.batch import DatasetBatch
 from dataclasses import dataclass
 from ml_gym.blueprints.component_factory import ComponentFactory, Injector
diff --git a/example/legacy/cross_validation/train_cv_mnist_conv_net.sh b/example/legacy/cross_validation/train_cv_mnist_conv_net.sh
index 5be310cb..81ff6a43 100644
--- a/example/legacy/cross_validation/train_cv_mnist_conv_net.sh
+++ b/example/legacy/cross_validation/train_cv_mnist_conv_net.sh
@@ -5,7 +5,6 @@
 python ../run.py --validation_mode CROSS_VALIDATION \
               --process_count 2 \
               --dashify_logging_path ../../dashify_logging/ \
-              --text_logging_path ../../general_logging/ \
               --gs_config_path gs_config_cv.yml \
               --evaluation_config_path cv_config.yml \
               --num_epochs 3 \
diff --git a/example/legacy/run.py b/example/legacy/run.py
index 6a97bbed..426d3687 100644
--- a/example/legacy/run.py
+++ b/example/legacy/run.py
@@ -6,7 +6,7 @@
 from ml_gym.starter import mlgym_entry_train, mlgym_entry_warm_start
 from ml_gym.validation.validator_factory import get_validator
 from ml_gym.io.config_parser import YAMLConfigLoader
-from typing import List, Dict, Type
+from typing import List
 from ml_gym.persistency.io import GridSearchAPIClientConfig, GridSearchAPIClientConstructableIF, GridSearchAPIClientConstructable, \
     GridSearchAPIClientType
 from pathlib import Path
@@ -54,10 +54,8 @@ def entry_train(args):
                       gs_config_raw_string=gs_config_string,
                       validation_strategy_config_raw_string=validation_strategy_config_string,
                       validator=validator,
-                      text_logging_path=args.text_logging_path,
                       process_count=args.process_count,
                       gpus=args.gpus,
-                      log_std_to_file=args.log_std_to_file,
                       num_epochs=args.num_epochs)
 
 
@@ -70,10 +68,8 @@ def entry_warm_start(args):
                            grid_search_id=args.grid_search_id,
                            logger_collection_constructable=logger_collection_constructable,
                            gs_api_client_constructable=gs_restful_api_client_constructable,
-                           text_logging_path=args.text_logging_path,
                            process_count=args.process_count,
                            gpus=args.gpus,
-                           log_std_to_file=args.log_std_to_file,
                            num_epochs=args.num_epochs)
 
     # starter = MLGymTrainWarmStarter(grid_search_id=grid_search_id, blueprint_class=blueprint_class, validation_mode=)
@@ -102,9 +98,7 @@ def parse_args_and_run():
     # Common
     parser.add_argument('--num_epochs', type=int, required=True, help='Number of epoch')
     parser.add_argument('--process_count', type=int, required=True, help='Max. number of processes running at a time.')
-    parser.add_argument('--text_logging_path', type=str, required=True, help='Path to python textual logging directory')
     parser.add_argument('--gpus', type=int, nargs='+', help='Indices of GPUs to distribute the GS over', default=None)
-    parser.add_argument('--log_std_to_file', default=False, action="store_true", help='Flag for forwarding std output to file')
     parser.add_argument('--websocket_logging_servers', type=str, nargs='+',
                         help='List of websocket logging servers, e.g., 127.0.0.1:9090 127.0.0.1:8080', default=None)
     parser.add_argument('--gs_rest_api_endpoint', type=str, required=True,
diff --git a/example/legacy/train_gs_mnist_conv_net.sh b/example/legacy/train_gs_mnist_conv_net.sh
index 80786921..dc2c42c7 100644
--- a/example/legacy/train_gs_mnist_conv_net.sh
+++ b/example/legacy/train_gs_mnist_conv_net.sh
@@ -5,7 +5,6 @@
 python run.py --validation_mode GRID_SEARCH \
               --process_count 5 \
               --dashify_logging_path dashify_logging/ \
-              --text_logging_path general_logging/ \
               --gs_config_path grid_search/gs_config.yml \
               --num_epochs 20 \
               --gpus 0 1 2 3 \
diff --git a/example/legacy/train_gs_mnist_conv_net_websocket.sh b/example/legacy/train_gs_mnist_conv_net_websocket.sh
index 3645c1bf..33b51638 100644
--- a/example/legacy/train_gs_mnist_conv_net_websocket.sh
+++ b/example/legacy/train_gs_mnist_conv_net_websocket.sh
@@ -5,7 +5,6 @@
 python run.py --validation_mode GRID_SEARCH \
               --run_mode TRAIN \
               --process_count 1 \
-              --text_logging_path general_logging/ \
               --gs_config_path grid_search/gs_config.yml \
               --num_epochs 20 \
               --gpus 0 1 2 3 \
diff --git a/src/ml_board/README.md b/src/ml_board/README.md
index ffc8f726..fd2282f8 100644
--- a/src/ml_board/README.md
+++ b/src/ml_board/README.md
@@ -74,6 +74,7 @@ payload:
 
 _GET /grid_searches/<grid_search_id>/experiments_
 
+TODO need to check "experiments" in URL  
 ```json
 [
     {
@@ -138,10 +139,17 @@ payload:
 }
 ```
 
-Delete Checkpoint Resources:
+Delete Checkpoint Resources: _model, optimiyer, stateful_component and lr_scheduler_ for a checkpoint ID in the experiment  :
+
+_DELETE /checkpoints/<grid_search_id>/<experiment_id>/<checkpoint_id>_
+
+Delete specific Checkpoint Resource:
+
+_DELETE /checkpoints/<grid_search_id><experiment_id>/<checkpoint_id>_
 
 _DELETE /checkpoints/<grid_search_id>/<experiment_id>/<checkpoint_id>/model_
 
+
 _DELETE /checkpoints/<grid_search_id>/<experiment_id>/<checkpoint_id>/optimizer_
 
 _DELETE /checkpoints/<grid_search_id>/<experiment_id>/<checkpoint_id>/stateful_component_
@@ -266,9 +274,6 @@ metric scores of a model at a specific epoch.
 }
 ```
 
-The files received by the websocket server are stored in
-`event_storage/mlgym_event_subscribers/<grid_search_id/event_storage_id><experiment_id>/checkpointing/<checkpoint_id>` as `model.pt`, `optimizer.pt` and `stateful_components.pt`. These files are available via the RESTful API to the clients and are not streamed.
-
 # Implementation idea:
 
 Add subscribers to trainer, evaluator and gymjob. We need to inject them via
diff --git a/src/ml_board/backend/restful_api/restful_api_server.py b/src/ml_board/backend/restful_api/restful_api_server.py
index 709dcaa4..f34bc333 100644
--- a/src/ml_board/backend/restful_api/restful_api_server.py
+++ b/src/ml_board/backend/restful_api/restful_api_server.py
@@ -45,6 +45,11 @@ def __init__(self, data_access: DataAccessIF):
             methods=["POST"],
             endpoint=self.add_checkpoint_resource,
         )
+        self.app.add_api_route(
+            path="/checkpoints/{grid_search_id}/{experiment_id}/{epoch}",
+            methods=["DELETE"],
+            endpoint=self.delete_checkpoints,
+        )
         self.app.add_api_route(
             path="/checkpoints/{grid_search_id}/{experiment_id}/{epoch}/{checkpoint_resource}",
             methods=["DELETE"],
@@ -198,6 +203,32 @@ def delete_checkpoint_resource(self, grid_search_id: str, experiment_id: str, ep
         except InvalidPathError as e:
             raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Provided invalid parameters for checkpoint resource.") from e
 
+    def delete_checkpoints(self, grid_search_id: str, experiment_id: str, epoch: str):
+        """
+        ``HTTP DELETE`` Delete checkpoint resource pickle file
+          given the epoch, experiment ID & grid search ID.
+
+        :params:
+             grid_search_id (str): Grid Search ID
+             experiment_id (str): Experiment ID
+             epoch (str): Epoch number
+        """
+        try:
+            chekpoint_resources = [
+                CheckpointResource.model,
+                CheckpointResource.lr_scheduler,
+                CheckpointResource.optimizer,
+                CheckpointResource.stateful_components,
+            ]
+
+            for checkpoint in chekpoint_resources:
+                self.data_access.delete_checkpoint_resource(
+                    grid_search_id=grid_search_id, experiment_id=experiment_id, epoch=epoch, checkpoint_resource=checkpoint
+                )
+
+        except InvalidPathError as e:
+            raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Provided invalid parameters for checkpoint resource.") from e
+
     def add_checkpoint_resource(
         self, grid_search_id: str, experiment_id: str, epoch: str, checkpoint_resource: CheckpointResource, file: bytes = File(...)
     ):
diff --git a/src/ml_gym/accelerate/__init__.py b/src/ml_gym/accelerate/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/ml_gym/batching/batch.py b/src/ml_gym/batching/batch.py
index c4df44ea..58855e86 100644
--- a/src/ml_gym/batching/batch.py
+++ b/src/ml_gym/batching/batch.py
@@ -28,12 +28,13 @@ def traverse_apply(ds: Union[Dict, List, torch.Tensor], apply_fun: Callable[[tor
             return [TorchDeviceMixin.traverse_apply(d, apply_fun) for d in ds]
         return apply_fun(ds)
 
+    @property
     @abstractmethod
-    def get_device(self) -> torch.device:
+    def device(self) -> torch.device:
         raise NotImplementedError
 
     @abstractmethod
-    def to_device(self, device: torch.device):
+    def to(self, device: torch.device) -> "Batch":
         raise NotImplementedError
 
     @abstractmethod
@@ -119,15 +120,18 @@ def detach(self):
         self._tags = self._tags.detach()
         self._samples = self._samples.detach()
 
-    def to_device(self, device: torch.device):
+    def to(self, device: torch.device) -> "DatasetBatch":
         self._samples = self._samples.to(device)
         self._targets = {k: v.to(device) for k, v in self._targets.items()}
         self._tags = self._tags.to(device)
+        return self
 
-    def to_cpu(self):
-        self.to_device(device=torch.device("cpu"))
+    def to_cpu(self) -> "DatasetBatch":
+        self.to(device=torch.device("cpu"))
+        return self
 
-    def get_device(self) -> torch.device:
+    @property
+    def device(self) -> torch.device:
         return self._samples.device
 
     @staticmethod
@@ -162,19 +166,23 @@ class InferenceResultBatch(Batch, TorchDeviceMixin):
 
     def __init__(self, targets: Dict[str, torch.Tensor] = None, predictions: Dict[str, torch.Tensor] = None, tags: torch.Tensor = None):
         self._targets = targets if targets is not None else {}
-        self._tags = tags
+        self._tags = tags if tags is not None else torch.tensor()
         self._predictions = predictions if predictions is not None else {}
+        self.to(self.device)
 
     def to_cpu(self):
-        self.to_device(device=torch.device("cpu"))
+        self.to(device=torch.device("cpu"))
 
-    def get_device(self) -> torch.device:
-        return self._tags.device
+    @property
+    def device(self) -> torch.device:
+        key = list(self._targets.keys())[0]
+        return self._targets[key].device
 
-    def to_device(self, device: torch.device):
+    def to(self, device: torch.device) -> "InferenceResultBatch":
         self._predictions = TorchDeviceMixin._dict_tensor_to_device(self._predictions, device)
         self._targets = {k: v.to(device) for k, v in self._targets.items()}
         self._tags = self._tags.to(device)
+        return self
 
     def detach(self):
         self._targets = {k: v.detach() for k, v in self._targets.items()}
@@ -263,7 +271,7 @@ def combine_pair(b_1: 'InferenceResultBatch', b_2: 'InferenceResultBatch') -> 'I
 
     @staticmethod
     def combine_impl(batches: List['InferenceResultBatch']) -> 'InferenceResultBatch':
-        tags = [tag for batch in batches for tag in batch.tags]
+        tags = torch.tensor([tag for batch in batches for tag in batch.tags])
         predictions = Batch._combine_tensor_dicts([batch.predictions for batch in batches])
         targets = Batch._combine_tensor_dicts([batch.targets for batch in batches])
         return InferenceResultBatch(targets=targets, predictions=predictions, tags=tags)
diff --git a/src/ml_gym/blueprints/blue_prints.py b/src/ml_gym/blueprints/blue_prints.py
index 9276f3e8..a07e1510 100644
--- a/src/ml_gym/blueprints/blue_prints.py
+++ b/src/ml_gym/blueprints/blue_prints.py
@@ -1,36 +1,26 @@
 from abc import ABC, abstractmethod
-from ml_gym.gym.jobs import GymJob
 from typing import List, Type, Dict, Any
 from ml_gym.modes import RunMode
-from ml_gym.persistency.logging import MLgymStatusLoggerCollectionConstructable
 import torch
-from ml_gym.persistency.io import GridSearchAPIClientConstructableIF
 
 
 class BluePrint(ABC):
-    """ Abstract class that provides a blueprint for creating `AbstractGymJob`
+    """ Abstract class that provides a blueprint for creating all the components for the GymJob`
     """
 
-    def __init__(self, run_mode: RunMode, num_epochs: int,
-                 config: Dict[str, Any], grid_search_id: str,
-                 gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-                 experiment_id: str,
-                 external_injection: Dict[str, Any] = None,
-                 logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None,
+    def __init__(self, run_mode: RunMode, config: Dict[str, Any], grid_search_id: str,
+                 experiment_id: str, external_injection: Dict[str, Any] = None,
                  warm_start_epoch: int = 0):
 
         self.run_mode = run_mode
         self.config = config
         self.grid_search_id = grid_search_id
         self.experiment_id = experiment_id
-        self.num_epochs = num_epochs
         self.external_injection = external_injection if external_injection is not None else {}
-        self.logger_collection_constructable = logger_collection_constructable
         self.warm_start_epoch = warm_start_epoch
-        self.gs_api_client_constructable = gs_api_client_constructable
 
     @abstractmethod
-    def construct(self, device: torch.device = None) -> GymJob:
+    def construct(self, device: torch.device = None) -> Dict:
         raise NotImplementedError
 
     @staticmethod
@@ -44,20 +34,14 @@ def create_blueprint(blue_print_class: Type["BluePrint"],
                          run_mode: RunMode,
                          experiment_config: Dict[str, Any],
                          experiment_id: str,
-                         num_epochs: int,
                          grid_search_id: str,
-                         gs_api_client_constructable: GridSearchAPIClientConstructableIF,
                          external_injection: Dict[str, Any] = None,
-                         logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None,
                          warm_start_epoch: int = 0) -> List["BluePrint"]:
 
         blue_print = blue_print_class(grid_search_id=grid_search_id,
                                       experiment_id=experiment_id,
-                                      num_epochs=num_epochs,
                                       warm_start_epoch=warm_start_epoch,
                                       run_mode=run_mode,
                                       config=experiment_config,
-                                      external_injection=external_injection,
-                                      logger_collection_constructable=logger_collection_constructable,
-                                      gs_api_client_constructable=gs_api_client_constructable)
+                                      external_injection=external_injection)
         return blue_print
diff --git a/src/ml_gym/blueprints/component_factory.py b/src/ml_gym/blueprints/component_factory.py
index fffae9ed..67499833 100644
--- a/src/ml_gym/blueprints/component_factory.py
+++ b/src/ml_gym/blueprints/component_factory.py
@@ -3,15 +3,16 @@
 from collections import namedtuple
 from dataclasses import dataclass, field
 from ml_gym.error_handling.exception import ComponentConstructionError, InjectMappingNotFoundError, DependentComponentNotFoundError
-from ml_gym.blueprints.constructables import ComponentConstructable, DatasetIteratorConstructable, \
-    DatasetIteratorSplitsConstructable, DeprecatedDataLoadersConstructable, EarlyStoppingRegistryConstructable, EarlyStoppingStrategyConstructable, LRSchedulerConstructable, OptimizerBundleConstructable, Requirement, DataLoadersConstructable, DatasetRepositoryConstructable, \
+from ml_gym.blueprints.constructables import AccelerateEvalComponentConstructable, AccelerateEvaluatorConstructable, AccelerateTrainComponentConstructable, AccelerateTrainerConstructable, ComponentConstructable, DatasetIteratorConstructable, \
+    DatasetIteratorSplitsConstructable, DeprecatedDataLoadersConstructable, EarlyStoppingRegistryConstructable, \
+    EarlyStoppingStrategyConstructable, LRSchedulerConstructable, OptimizerBundleConstructable, \
+    Requirement, DataLoadersConstructable, DatasetRepositoryConstructable, \
     OptimizerConstructable, ModelRegistryConstructable, ModelConstructable, LossFunctionRegistryConstructable, \
     MetricFunctionRegistryConstructable, TrainerConstructable, EvaluatorConstructable, MappedLabelsIteratorConstructable, \
     FilteredLabelsIteratorConstructable, FeatureEncodedIteratorConstructable, CombinedDatasetIteratorConstructable, \
     DataCollatorConstructable, PredictionPostProcessingRegistryConstructable, TrainComponentConstructable, EvalComponentConstructable, \
     IteratorViewConstructable, OneHotEncodedTargetsIteratorConstructable, InMemoryDatasetIteratorConstructable, \
     ShuffledDatasetIteratorConstructable, CheckpointingStrategyConstructable, CheckpointingRegistryConstructable
-# from ml_gym.util.logger import LogLevel, ConsoleLogger
 
 
 class Injector:
@@ -132,9 +133,13 @@ def __init__(self, injector: Injector = None):
             ComponentVariant("PREDICTION_POSTPROCESSING_REGISTRY", "DEFAULT", PredictionPostProcessingRegistryConstructable),
             ComponentVariant("MODEL", "DEFAULT", ModelConstructable),
             ComponentVariant("TRAIN_COMPONENT", "DEFAULT", TrainComponentConstructable),
+            ComponentVariant("TRAIN_COMPONENT", "ACCELERATE", AccelerateTrainComponentConstructable),
             ComponentVariant("TRAINER", "DEFAULT", TrainerConstructable),
+            ComponentVariant("TRAINER", "ACCELERATE", AccelerateTrainerConstructable),
             ComponentVariant("EVAL_COMPONENT", "DEFAULT", EvalComponentConstructable),
             ComponentVariant("EVALUATOR", "DEFAULT", EvaluatorConstructable),
+            ComponentVariant("EVAL_COMPONENT", "ACCELERATE", AccelerateEvalComponentConstructable),
+            ComponentVariant("EVALUATOR", "ACCELERATE", AccelerateEvaluatorConstructable),
             ComponentVariant("EARLY_STOPPING_STRATEGY_REGISTRY", "DEFAULT", EarlyStoppingRegistryConstructable),
             ComponentVariant("EARLY_STOPPING_STRATEGY", "DEFAULT", EarlyStoppingStrategyConstructable),
             ComponentVariant("CHECKPOINTING_STRATEGY_REGISTRY", "DEFAULT", CheckpointingRegistryConstructable),
diff --git a/src/ml_gym/blueprints/constructables.py b/src/ml_gym/blueprints/constructables.py
index 41b83ebd..42ea6410 100644
--- a/src/ml_gym/blueprints/constructables.py
+++ b/src/ml_gym/blueprints/constructables.py
@@ -8,6 +8,8 @@
 from data_stack.mnist.factory import MNISTFactory
 from ml_gym.data_handling.dataset_loader import DatasetLoader, DatasetLoaderFactory
 from ml_gym.early_stopping.early_stopping_strategies import EarlyStoppingIF, EarlyStoppingStrategyFactory
+from ml_gym.gym.evaluators.accelerate_evaluator import AccelerateEvalComponent, AccelerateEvaluator
+from ml_gym.gym.trainers.accelerate_trainer import AccelerateTrainComponent, AccelerateTrainer
 from ml_gym.optimizers.lr_scheduler_factory import LRSchedulerFactory
 from ml_gym.optimizers.lr_schedulers import LRSchedulerAdapter
 from ml_gym.optimizers.optimizer import OptimizerAdapter, OptimizerBundle
@@ -15,12 +17,12 @@
 from ml_gym.models.nn.net import NNModel
 from collections.abc import Mapping
 from ml_gym.registries.class_registry import ClassRegistry
-from ml_gym.gym.trainer import Trainer, TrainComponent, InferenceComponent
+from ml_gym.gym.trainers.standard_trainer import Trainer, TrainComponent, InferenceComponent
 from ml_gym.loss_functions.loss_functions import Loss
 from sklearn.metrics import f1_score, recall_score, precision_score, accuracy_score, balanced_accuracy_score
 from ml_gym.metrics.metrics import Metric, binary_aupr_score, binary_auroc_score
 from ml_gym.metrics.metric_factory import MetricFactory
-from ml_gym.gym.evaluator import Evaluator, EvalComponent
+from ml_gym.gym.evaluators.evaluator import Evaluator, EvalComponent
 from ml_gym.data_handling.postprocessors.factory import ModelGymInformedIteratorFactory
 from ml_gym.data_handling.postprocessors.collator import Collator
 from ml_gym.gym.post_processing import PredictPostProcessingIF, SoftmaxPostProcessorImpl, \
@@ -431,12 +433,13 @@ def _construct_impl(self):
 class ModelConstructable(ComponentConstructable):
     model_type: str = ""
     model_definition: Dict[str, Any] = field(default_factory=dict)
-    seed: int = 0
+    seed: int = None
     prediction_publication_keys: Dict[str, str] = field(default_factory=dict)
 
     def _construct_impl(self) -> NNModel:
+        other_params = {"seed": self.seed} if self.seed is not None else {}
         model_type = self.get_requirement("model_registry")
-        return model_type(seed=self.seed, **self.model_definition, **self.prediction_publication_keys)
+        return model_type(**other_params, **self.model_definition, **self.prediction_publication_keys)
 
 
 @dataclass
@@ -453,7 +456,24 @@ def _construct_impl(self) -> TrainComponent:
                           for config in self.post_processors_config]
 
         inference_component = InferenceComponent(no_grad=False)
-        train_component = TrainComponent(inference_component, postprocessors, train_loss_fun, self.show_progress)
+        train_component = TrainComponent(inference_component, postprocessors, train_loss_fun)
+        return train_component
+
+
+@dataclass
+class AccelerateTrainComponentConstructable(ComponentConstructable):
+    loss_fun_config: Dict = field(default_factory=dict)
+    post_processors_config: List[Dict] = field(default_factory=list)
+
+    def _construct_impl(self) -> AccelerateTrainComponent:
+        prediction_post_processing_registry: ClassRegistry = self.get_requirement("prediction_postprocessing_registry")
+        loss_function_registry: ClassRegistry = self.get_requirement("loss_function_registry")
+        train_loss_fun = loss_function_registry.get_instance(**self.loss_fun_config)
+        postprocessors = [PredictPostProcessing(prediction_post_processing_registry.get_instance(config["key"], **config["params"]))
+                          for config in self.post_processors_config]
+
+        inference_component = InferenceComponent(no_grad=False)
+        train_component = AccelerateTrainComponent(inference_component, postprocessors, train_loss_fun)
         return train_component
 
 
@@ -463,13 +483,22 @@ class TrainerConstructable(ComponentConstructable):
     def _construct_impl(self) -> Trainer:
         train_loader: DatasetLoader = self.get_requirement("data_loaders")
         train_component: TrainComponent = self.get_requirement("train_component")
-        trainer = Trainer(train_component=train_component, train_loader=train_loader, verbose=True)
+        trainer = Trainer(train_component=train_component, train_loader=train_loader)
+        return trainer
+
+
+@dataclass
+class AccelerateTrainerConstructable(ComponentConstructable):
+
+    def _construct_impl(self) -> AccelerateTrainer:
+        train_loader: DatasetLoader = self.get_requirement("data_loaders")
+        train_component: TrainComponent = self.get_requirement("train_component")
+        trainer = AccelerateTrainer(train_component=train_component, train_loader=train_loader)
         return trainer
 
 
 @dataclass
 class EvalComponentConstructable(ComponentConstructable):
-    train_split_name: str = ""
     metrics_config: List = field(default_factory=list)
     loss_funs_config: List = field(default_factory=list)
     post_processors_config: List[Dict] = field(default_factory=list)
@@ -479,7 +508,7 @@ class EvalComponentConstructable(ComponentConstructable):
     metrics_computation_config: List[Dict] = None
     loss_computation_config: List[Dict] = None
 
-    def _construct_impl(self) -> Evaluator:
+    def _construct_impl(self) -> EvalComponent:
         dataset_loaders: Dict[str, DatasetLoader] = self.get_requirement("data_loaders")
         loss_function_registry: ClassRegistry = self.get_requirement("loss_function_registry")
         metric_registry: ClassRegistry = self.get_requirement("metric_registry")
@@ -502,7 +531,7 @@ def _construct_impl(self) -> Evaluator:
                     postprocessors_dict["default"].append(PredictPostProcessing(prediction_post_processing_registry.get_instance(**config)))
 
         inference_component = InferenceComponent(no_grad=True)
-        eval_component = EvalComponent(inference_component, postprocessors_dict, metric_funs, loss_funs, dataset_loaders, self.train_split_name,
+        eval_component = EvalComponent(inference_component, postprocessors_dict, metric_funs, loss_funs, dataset_loaders,
                                        self.show_progress, self.cpu_target_subscription_keys, self.cpu_prediction_subscription_keys,
                                        self.metrics_computation_config, self.loss_computation_config)
         return eval_component
@@ -517,6 +546,54 @@ def _construct_impl(self) -> Evaluator:
         return evaluator
 
 
+@dataclass
+class AccelerateEvalComponentConstructable(ComponentConstructable):
+    metrics_config: List = field(default_factory=list)
+    loss_funs_config: List = field(default_factory=list)
+    post_processors_config: List[Dict] = field(default_factory=list)
+    cpu_target_subscription_keys: List[str] = field(default_factory=list)
+    cpu_prediction_subscription_keys: List[str] = field(default_factory=list)
+    metrics_computation_config: List[Dict] = None
+    loss_computation_config: List[Dict] = None
+
+    def _construct_impl(self) -> AccelerateEvalComponent:
+        dataset_loaders: Dict[str, DatasetLoader] = self.get_requirement("data_loaders")
+        loss_function_registry: ClassRegistry = self.get_requirement("loss_function_registry")
+        metric_registry: ClassRegistry = self.get_requirement("metric_registry")
+        prediction_post_processing_registry: ClassRegistry = self.get_requirement("prediction_postprocessing_registry")
+
+        loss_funs = {conf["tag"]: loss_function_registry.get_instance(**conf) for conf in self.loss_funs_config}
+        metric_funs = [metric_registry.get_instance(**conf) for conf in self.metrics_config]
+
+        postprocessors_dict = defaultdict(list)
+        for config in self.post_processors_config:
+            if "applicable_splits" in config:
+                for split in config["applicable_splits"]:
+                    postprocessors_dict[split].append(PredictPostProcessing(
+                        prediction_post_processing_registry.get_instance(config["key"], **config["params"])))
+            else:
+                if "params" in config:
+                    postprocessors_dict["default"].append(PredictPostProcessing(
+                        prediction_post_processing_registry.get_instance(config["key"], **config["params"])))
+                else:  # TODO this is the legacy version!
+                    postprocessors_dict["default"].append(PredictPostProcessing(prediction_post_processing_registry.get_instance(**config)))
+
+        inference_component = InferenceComponent(no_grad=True)
+        eval_component = AccelerateEvalComponent(inference_component, postprocessors_dict, metric_funs, loss_funs, dataset_loaders,
+                                                 self.cpu_target_subscription_keys, self.cpu_prediction_subscription_keys,
+                                                 self.metrics_computation_config, self.loss_computation_config)
+        return eval_component
+
+
+@dataclass
+class AccelerateEvaluatorConstructable(ComponentConstructable):
+
+    def _construct_impl(self) -> AccelerateEvaluator:
+        eval_component: AccelerateEvalComponent = self.get_requirement("eval_component")
+        evaluator = AccelerateEvaluator(eval_component)
+        return evaluator
+
+
 @dataclass
 class EarlyStoppingRegistryConstructable(ComponentConstructable):
     class StrategyKeys:
diff --git a/src/ml_gym/cmd_entrypoint/__init__.py b/src/ml_gym/cmd_entrypoint/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/ml_gym/cmd_entrypoint/cmd.py b/src/ml_gym/cmd_entrypoint/cmd.py
new file mode 100644
index 00000000..726376ec
--- /dev/null
+++ b/src/ml_gym/cmd_entrypoint/cmd.py
@@ -0,0 +1,237 @@
+import argparse
+from dataclasses import dataclass
+from datetime import datetime
+from ml_board.backend.restful_api.data_models import FileFormat
+from ml_gym.blueprints.blue_prints import BluePrint
+from ml_gym.error_handling.exception import GymError
+from ml_gym.gym.gym import Gym, GymFactory
+from ml_gym.persistency.logging import MLgymStatusLoggerCollectionConstructable, LoggerConstructableIF, MLgymStatusLoggerConfig, \
+    MLgymStatusLoggerTypes
+from ml_gym.modes import RunMode, ValidationMode
+from ml_gym.validation.validator_factory import get_validator
+from ml_gym.io.config_parser import YAMLConfigLoader
+from typing import Callable, List, Tuple, Type, Dict, Union
+from ml_gym.persistency.io import GridSearchAPIClientConfig, GridSearchAPIClientConstructableIF, GridSearchAPIClientConstructable, \
+    GridSearchAPIClientType
+from pathlib import Path
+import yaml
+from accelerate import Accelerator
+import json
+
+
+@dataclass
+class MultiProcessingEnvironmentConfig:
+    process_count: int
+    computation_device_ids: List[int]
+
+
+@dataclass
+class MainProcessEnvironmentConfig:
+    computation_device_id: int
+
+
+@dataclass
+class AccelerateEnvironmentConfig:
+    pass
+
+
+@dataclass
+class TrainRunConfiguration:
+    num_epochs: int
+    gs_config_path: str
+    validation_config_path = None
+    num_batches_per_epoch: int = None
+
+
+@dataclass
+class WarmStartRunConfiguration:
+    gridsaerch_id: str
+    num_epochs: int
+    num_batches_per_epoch: int = None
+
+
+@dataclass
+class LoggingConfiguration:
+    websocket_logging_servers: List[str]
+    gs_rest_api_endpoint: str
+
+
+def get_gym_from_environment_config(env_config: Union[MultiProcessingEnvironmentConfig, MainProcessEnvironmentConfig,
+                                                      AccelerateEnvironmentConfig],
+                                    logger_collection_constructable: LoggerConstructableIF,
+                                    gs_restful_api_client_constructable: GridSearchAPIClientConstructableIF,
+                                    num_epochs: int,
+                                    num_batches_per_epoch: int = None,
+                                    accelerator: Accelerator = None) -> Gym:
+
+    if isinstance(env_config, MainProcessEnvironmentConfig):
+        gym = GymFactory.get_sequential_gym(logger_collection_constructable=logger_collection_constructable,
+                                            gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                                            device_id=env_config.computation_device_id,
+                                            num_batches_per_epoch=num_batches_per_epoch)
+    elif isinstance(env_config, AccelerateEnvironmentConfig):
+        gym = GymFactory.get_sequential_gym(logger_collection_constructable=logger_collection_constructable,
+                                            gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                                            num_epochs=num_epochs,
+                                            num_batches_per_epoch=num_batches_per_epoch,
+                                            accelerator=accelerator)
+    elif isinstance(env_config, MultiProcessingEnvironmentConfig):
+        gym = GymFactory.get_parallel_single_node_gym(logger_collection_constructable=logger_collection_constructable,
+                                                      gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                                                      process_count=env_config.process_count,
+                                                      num_epochs=num_epochs,
+                                                      device_ids=env_config.computation_device_ids,
+                                                      num_batches_per_epoch=num_batches_per_epoch)
+    else:
+        raise GymError("Did not provide correct env_config")
+
+    return gym
+
+
+def get_logger_constructable(websocket_logging_servers: List) -> LoggerConstructableIF:
+    websocket_logging_servers = [(f"{protocol}:{ip}", int(port)) for protocol, ip, port in [
+        connection.split(":") for connection in websocket_logging_servers]]
+
+    logging_configs = [MLgymStatusLoggerConfig(logger_type=MLgymStatusLoggerTypes.STREAMED_LOGGER, logger_config={"host": ip, "port": port})
+                       for ip, port in websocket_logging_servers]
+    logger_collection_constructable = MLgymStatusLoggerCollectionConstructable(logging_configs)
+    return logger_collection_constructable
+
+
+def get_grid_search_restful_api_client_constructable(endpoint: str) -> GridSearchAPIClientConstructableIF:
+    client_config = GridSearchAPIClientConfig(api_client_type=GridSearchAPIClientType.GRID_SEARCH_RESTFUL_API_CLIENT,
+                                              api_client_config={"endpoint": endpoint})
+    client_constructable = GridSearchAPIClientConstructable(client_config)
+    return client_constructable
+
+
+def entry_train(gridsearch_id: str, blueprint_class: Type[BluePrint], gym: Gym, gs_config_path: str,
+                validation_strategy_config_path: str, gs_restful_api_client_constructable: GridSearchAPIClientConstructableIF,
+                accelerator: Accelerator = None):
+
+    def log_configs(gs_config_string: str, validation_strategy_config_raw_string: str,
+                    gs_restful_api_client_constructable: GridSearchAPIClientConstructableIF, accelerator: Accelerator = None):
+        if accelerator is None or accelerator.is_main_process:
+            gs_api_client = gs_restful_api_client_constructable.construct()
+            gs_api_client.add_config_string(grid_search_id=gridsearch_id,
+                                            file_format=FileFormat.YAML,
+                                            config_name="grid_search_config.yml",
+                                            config=gs_config_string)
+
+            if validation_strategy_config_raw_string is not None:
+                gs_api_client.add_config_string(grid_search_id=gridsearch_id,
+                                                config_name="validation_strategy_config.yml",
+                                                config=validation_strategy_config_raw_string)
+
+            for blueprint in blueprints:
+                gs_api_client.add_config_string(grid_search_id=blueprint.grid_search_id, config_name="experiment_config.json",
+                                                config=json.dumps(blueprint.config), experiment_id=blueprint.experiment_id,
+                                                file_format=FileFormat.JSON)
+
+    gs_config_string = Path(gs_config_path).read_text()
+    gs_config = YAMLConfigLoader.load_string(gs_config_string)
+
+    if validation_strategy_config_path is not None:
+        validation_strategy_config_string = Path(validation_strategy_config_path).read_text()
+        validation_strategy_config = YAMLConfigLoader.load(validation_strategy_config_string)
+        validation_mode = ValidationMode[list(validation_strategy_config.keys())[0]]
+    else:
+        validation_strategy_config_string = None
+        validation_strategy_config = None
+        validation_mode = ValidationMode.GRID_SEARCH
+
+    gs_config = YAMLConfigLoader.load_string(gs_config_string)
+
+    validator = get_validator(validation_mode, blueprint_class, RunMode.TRAIN, validation_strategy_config, gs_config)
+
+    blueprints = validator.create_blueprints(grid_search_id=gridsearch_id,
+                                             blue_print_type=blueprint_class,
+                                             gs_config=gs_config)
+
+    log_configs(gs_config_string=gs_config_string, validation_strategy_config_raw_string=validation_strategy_config_string,
+                gs_restful_api_client_constructable=gs_restful_api_client_constructable, accelerator=accelerator)
+
+    gym.run(blueprints)
+
+
+def entry_warm_start(blueprint_class: Type[BluePrint], gym: Gym, grid_search_id: int,
+                     gs_restful_api_client_constructable: GridSearchAPIClientConstructableIF, num_epochs: int):
+
+    gs_api_client = gs_restful_api_client_constructable.construct()
+    experiment_statuses = gs_api_client.get_experiment_statuses(grid_search_id)
+
+    blueprints = [BluePrint.create_blueprint(blue_print_class=blueprint_class,
+                                             run_mode=RunMode.WARM_START,
+                                             experiment_config=experiment_status.experiment_config,
+                                             warm_start_epoch=experiment_status.last_checkpoint_id,
+                                             grid_search_id=grid_search_id,
+                                             experiment_id=experiment_status.experiment_id)
+                  for experiment_status in experiment_statuses if experiment_status.last_checkpoint_id < num_epochs]
+
+    gym.run(blueprints)
+
+
+def get_args() -> Tuple[Callable, Dict]:
+    parser = argparse.ArgumentParser(description='Run a grid search on CPUs or distributed over multiple GPUs')
+
+    parser.add_argument('--config_path', type=str, required=True, help='Path to the run configuration file')
+
+    args = parser.parse_args()
+    return args.config_path
+
+
+def parse_run_configuration(run_configuration_file_path: str) -> Tuple[Union[TrainRunConfiguration, WarmStartRunConfiguration],
+                                                                       Union[MultiProcessingEnvironmentConfig, AccelerateEnvironmentConfig,
+                                                                       MainProcessEnvironmentConfig], LoggingConfiguration]:
+    with open(run_configuration_file_path, "r") as fp:
+        run_configuration_dict = yaml.safe_load(fp)
+
+    run_dict = run_configuration_dict["run_configuration"]
+    environment_dict = run_configuration_dict["environment"]
+    logging_dict = run_configuration_dict["logging"]
+
+    if run_dict["type"] == "train":
+        run_config = TrainRunConfiguration(**run_dict["config"])
+    else:
+        run_config = WarmStartRunConfiguration(**run_dict["config"])
+
+    if environment_dict["type"] == "multiprocessing":
+        environment_config = MultiProcessingEnvironmentConfig(**environment_dict["config"])
+    elif environment_dict["type"] == "accelerate":
+        environment_config = AccelerateEnvironmentConfig()
+    else:
+        environment_config = MainProcessEnvironmentConfig(**environment_dict["config"])
+
+    logging_config = LoggingConfiguration(**logging_dict)
+
+    return run_config, environment_config, logging_config
+
+
+def get_logging_constructables(logging_config: LoggingConfiguration) -> Tuple[LoggerConstructableIF, GridSearchAPIClientConstructableIF]:
+    # get logger constructables
+    logger_collection_constructable = get_logger_constructable(logging_config.websocket_logging_servers)
+    gs_restful_api_client_constructable = get_grid_search_restful_api_client_constructable(endpoint=logging_config.gs_rest_api_endpoint)
+    return logger_collection_constructable, gs_restful_api_client_constructable
+
+
+def run(blueprint_class: BluePrint, run_configuration_file_path):
+
+    run_config, env_config, logging_config = parse_run_configuration(run_configuration_file_path=run_configuration_file_path)
+
+    logger_collection_constructable, gs_restful_api_client_constructable = get_logging_constructables(logging_config)
+    
+    accelerator: Accelerator = Accelerator() if isinstance(env_config, AccelerateEnvironmentConfig) else None
+
+    gym = get_gym_from_environment_config(env_config, logger_collection_constructable, gs_restful_api_client_constructable,
+                                          run_config.num_epochs, run_config.num_batches_per_epoch, accelerator)
+
+
+    if isinstance(run_config, TrainRunConfiguration):
+        gridsearch_id = datetime.now().strftime("%Y-%m-%d--%H-%M-%S")
+        entry_train(gridsearch_id=gridsearch_id, blueprint_class=blueprint_class, gym=gym, gs_config_path=run_config.gs_config_path,
+                    validation_strategy_config_path=run_config.validation_config_path,
+                    gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                    accelerator=accelerator)
+    elif isinstance(run_config, WarmStartRunConfiguration):
+        entry_warm_start(blueprint_class=blueprint_class, gym=gym, grid_search_id=run_config.gridsaerch_id,
+                         gs_restful_api_client_constructable=gs_restful_api_client_constructable, num_epochs=run_config.num_epochs)
diff --git a/src/ml_gym/data_handling/dataset_loader.py b/src/ml_gym/data_handling/dataset_loader.py
index f4337bd1..78fa0a79 100644
--- a/src/ml_gym/data_handling/dataset_loader.py
+++ b/src/ml_gym/data_handling/dataset_loader.py
@@ -1,3 +1,5 @@
+from abc import ABC, abstractmethod
+from itertools import chain
 from ml_gym.error_handling.exception import SamplerNotFoundError
 from torch.utils.data import DataLoader
 from torch.utils.data.sampler import RandomSampler, WeightedRandomSampler, Sampler, SequentialSampler
@@ -7,6 +9,7 @@
 import torch
 from ml_gym.data_handling.postprocessors.collator import Collator
 from enum import Enum
+from accelerate.data_loader import DataLoaderShard
 
 
 class DatasetLoaderFactory:
@@ -58,6 +61,12 @@ def get_splitted_data_loaders(dataset_splits: Dict[str, InformedDatasetIteratorI
                                                      drop_last=drop_last)
         return data_loaders
 
+    @staticmethod
+    def get_data_loader_shard_wrapper(data_loader_shard: DataLoaderShard, dataset_name: str,
+                                      dataset_tag: str) -> "DataLoaderShardWrapper":
+        data_loader = DataLoaderShardWrapper(data_loader_shard=data_loader_shard, dataset_name=dataset_name, dataset_tag=dataset_tag)
+        return data_loader
+
 
 class SamplerFactory:
 
@@ -96,7 +105,20 @@ def get_sequential_sampler(dataset: InformedDatasetIteratorIF) -> Sampler:
         return SequentialSampler(data_source=dataset)
 
 
-class DatasetLoader(DataLoader):
+class DatsetLoaderMetaInfoIF(ABC):
+
+    @property
+    @abstractmethod
+    def dataset_name(self) -> str:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def dataset_tag(self) -> str:
+        raise NotImplementedError
+
+
+class DatasetLoader(DatsetLoaderMetaInfoIF, DataLoader):
     def __init__(self, dataset_iterator: InformedDatasetIteratorIF, batch_size: int, sampler: Sampler,
                  collate_fn: Collator = None, drop_last: bool = False):
         super().__init__(dataset=dataset_iterator, sampler=sampler, batch_size=batch_size, collate_fn=collate_fn, drop_last=drop_last)
@@ -117,3 +139,24 @@ def device(self) -> torch.device:
     def device(self, d: torch.device):
         if self.collate_fn is not None:
             self.collate_fn.device = d
+
+
+class DataLoaderShardWrapper(DatsetLoaderMetaInfoIF, DataLoader):
+    def __init__(self, data_loader_shard: DataLoaderShard, dataset_name: str, dataset_tag: str):
+        self._dataset_name = dataset_name
+        self._dataset_tag = dataset_tag
+        self.data_loader_shard = data_loader_shard
+
+    def __iter__(self):
+        return iter(self.data_loader_shard)
+
+    def __len__(self) -> int:
+        return len(self.data_loader_shard)
+
+    @property
+    def dataset_name(self) -> str:
+        return self._dataset_name
+
+    @property
+    def dataset_tag(self) -> str:
+        return self._dataset_tag
diff --git a/src/ml_gym/error_handling/exception.py b/src/ml_gym/error_handling/exception.py
index 7294bc55..d661f954 100644
--- a/src/ml_gym/error_handling/exception.py
+++ b/src/ml_gym/error_handling/exception.py
@@ -67,11 +67,6 @@ class ClassRegistryItemInstantiationError(Exception):
     pass
 
 
-class SingletonAlreadyInstantiatedError(Exception):
-    """Exception raised when trying to get another instance from a singleton class"""
-    pass
-
-
 class InjectMappingNotFoundError(Exception):
     """Exception raised when Injector cannot inject given placehold when it was not defined in the mapping."""
     pass
@@ -129,3 +124,13 @@ class InvalidPathError(Exception):
 class CheckpointEntityError(Exception):
     """Raised when there is an error within the checkpoint entity."""
     pass
+
+
+class EarlyStoppingCriterionFulfilledError(Exception):
+    """Raised when the early stopping criterion in the gym job was fulfilled."""
+    pass
+
+
+class GymError(Exception):
+    """Raised when an error occurs within the gym or during gym instantiation."""
+    pass
diff --git a/src/ml_gym/gym/evaluators/__init__.py b/src/ml_gym/gym/evaluators/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/ml_gym/gym/evaluators/accelerate_evaluator.py b/src/ml_gym/gym/evaluators/accelerate_evaluator.py
new file mode 100644
index 00000000..4136f1f2
--- /dev/null
+++ b/src/ml_gym/gym/evaluators/accelerate_evaluator.py
@@ -0,0 +1,163 @@
+from enum import Enum
+from typing import Any, Callable, Dict, List, Union
+from ml_gym.gym.evaluators.evaluator import AbstractEvaluator
+from ml_gym.gym.post_processing import PredictPostProcessingIF
+from ml_gym.persistency.logging import ExperimentStatusLogger
+import torch
+from ml_gym.batching.batch import DatasetBatch, EvaluationBatchResult, InferenceResultBatch
+from ml_gym.data_handling.dataset_loader import DatasetLoader
+from ml_gym.gym.inference_component import InferenceComponent
+from ml_gym.metrics.metrics import Metric
+from ml_gym.models.nn.net import NNModel
+from ml_gym.loss_functions.loss_functions import Loss
+from ml_gym.gym.predict_postprocessing_component import PredictPostprocessingComponent
+from ml_gym.error_handling.exception import BatchStateError, EvaluationError, MetricCalculationError, LossCalculationError
+from accelerate import Accelerator
+
+
+class AccelerateEvaluator(AbstractEvaluator):
+    def __init__(self, eval_component: "AccelerateEvalComponent"):
+        self.eval_component = eval_component
+
+    def evaluate(self, model: NNModel, accelerator: Accelerator, epoch_result_callback_fun: Callable = None,
+                 batch_processed_callback_fun: Callable = None) -> List[EvaluationBatchResult]:
+        model = model.eval()
+
+        # returns a EvaluationBatchResult for each split
+        evaluation_batch_results = self.eval_component.evaluate(model=model, accelerator=accelerator,
+                                                                epoch_result_callback_fun=epoch_result_callback_fun,
+                                                                batch_processed_callback_fun=batch_processed_callback_fun)
+        return evaluation_batch_results
+
+
+class AccelerateEvalComponent:
+    """This thing always comes with batteries included, i.e., datasets, loss functions etc. are all already stored in here."""
+
+    def __init__(self, inference_component: InferenceComponent, post_processors: Dict[str, PredictPostprocessingComponent], metrics: List[Metric],
+                 loss_funs: Dict[str, Loss], dataset_loaders: Dict[str, DatasetLoader],
+                 cpu_target_subscription_keys: List[str] = None, cpu_prediction_subscription_keys: List[Union[str, List]] = None,
+                 metrics_computation_config: List[Dict] = None, loss_computation_config: List[Dict] = None):
+        self.loss_funs = loss_funs
+        self.inference_component = inference_component
+        # maps split names to postprocessors
+        self.post_processors = post_processors
+        self.metrics = metrics
+        self.dataset_loaders = dataset_loaders
+        self.cpu_target_subscription_keys = cpu_target_subscription_keys
+        self.cpu_prediction_subscription_keys = cpu_prediction_subscription_keys
+        # determines which metrics are applied to which splits (metric_key to split list)
+        self.metrics_computation_config = None if metrics_computation_config is None else {
+            m["metric_tag"]: m["applicable_splits"] for m in metrics_computation_config}
+        # determines which losses are applied to which splits (loss_key to split list)
+        self.loss_computation_config = None if loss_computation_config is None else {
+            m["loss_tag"]: m["applicable_splits"] for m in loss_computation_config}
+        self.experiment_status_logger: ExperimentStatusLogger = None
+
+    def evaluate(self, model: NNModel, accelerator: Accelerator, epoch_result_callback_fun: Callable = None,
+                 batch_processed_callback_fun: Callable = None) -> List[EvaluationBatchResult]:
+        return [self.evaluate_dataset_split(model, split_name, loader, accelerator, epoch_result_callback_fun, batch_processed_callback_fun) for split_name, loader in self.dataset_loaders.items()]
+
+    def evaluate_dataset_split(self, model: NNModel, split_name: str,
+                               dataset_loader: DatasetLoader, accelerator: Accelerator, epoch_result_callback_fun: Callable = None,
+                               batch_processed_callback_fun: Callable = None) -> EvaluationBatchResult:
+        post_processors = self.post_processors[split_name] + self.post_processors["default"]
+
+        # calc losses
+        if self.loss_computation_config is not None:
+            loss_tags = [loss_tag for loss_tag, applicable_splits in self.loss_computation_config.items()
+                         if split_name in applicable_splits]
+            split_loss_funs = {tag: loss_fun for tag, loss_fun in self.loss_funs.items() if tag in loss_tags}
+        else:
+            split_loss_funs = self.loss_funs
+
+        batch_losses = []
+        inference_result_batches_cpu = []
+        num_batches = len(dataset_loader)
+        processed_batches = 0
+        for batch in dataset_loader:
+            inference_result_batch = self.forward_batch(dataset_batch=batch, model=model, postprocessors=post_processors)
+
+            irb_filtered_dict = {"predictions": inference_result_batch.predictions, "targets": inference_result_batch.targets,
+                                 "tags": inference_result_batch.tags}
+            irb_filtered_dict_gathered = accelerator.gather_for_metrics(irb_filtered_dict)
+
+            irb_filtered_gathered = InferenceResultBatch(predictions=irb_filtered_dict_gathered["predictions"],
+                                                         targets=irb_filtered_dict_gathered["targets"],
+                                                         tags=irb_filtered_dict_gathered["tags"])
+
+            batch_loss = self._calculate_loss_scores(irb_filtered_gathered, split_loss_funs)
+            batch_losses.append(batch_loss)
+
+            irb_filtered = irb_filtered_gathered.split_results(predictions_keys=self.cpu_prediction_subscription_keys,
+                                                               target_keys=self.cpu_target_subscription_keys,
+                                                               device=torch.device("cpu"))
+
+            inference_result_batches_cpu.append(irb_filtered)
+            processed_batches += 1
+            splits = list(self.dataset_loaders.keys())
+            if accelerator.is_main_process:
+                batch_processed_callback_fun(status="evaluation",
+                                             num_batches=num_batches,
+                                             current_batch=processed_batches,
+                                             splits=splits,
+                                             current_split=split_name)
+
+        # calc metrics
+        try:
+            prediction_batch = InferenceResultBatch.combine(inference_result_batches_cpu)
+        except BatchStateError as e:
+            raise EvaluationError(f"Error combining inference result batch on split {split_name}.") from e
+
+        # select metrics for split
+        if self.metrics_computation_config is not None:
+            metric_tags = [metric_tag for metric_tag, applicable_splits in self.metrics_computation_config.items()
+                           if split_name in applicable_splits]
+            split_metrics = [metric for metric in self.metrics if metric.tag in metric_tags]
+        else:
+            split_metrics = self.metrics
+        metric_scores = self._calculate_metric_scores(prediction_batch, split_metrics)
+
+        # aggregate losses
+        loss_keys = batch_losses[0].keys()
+        loss_scores = {key: [torch.mean(torch.Tensor([l[key] for l in batch_losses])).item()] for key in loss_keys}
+
+        evaluation_result = EvaluationBatchResult(losses=loss_scores,
+                                                  metrics=metric_scores,
+                                                  dataset_name=dataset_loader.dataset_name,
+                                                  split_name=split_name)
+        if epoch_result_callback_fun is not None and accelerator.is_main_process:
+            epoch_result_callback_fun(evaluation_result=evaluation_result)
+        return evaluation_result
+
+    def _get_metric_fun(self, identifier: str, target_subscription: Enum, prediction_subscription: Enum,
+                        metric_fun: Callable, params: Dict[str, Any]) -> Metric:
+        return Metric(identifier, target_subscription, prediction_subscription, metric_fun, params)
+
+    def forward_batch(self, dataset_batch: DatasetBatch, model: NNModel,
+                      postprocessors: List[PredictPostProcessingIF]) -> InferenceResultBatch:
+        inference_result_batch = self.inference_component.predict(model, dataset_batch, postprocessors)
+        return inference_result_batch
+
+    def _calculate_metric_scores(self, inference_batch: InferenceResultBatch, split_metrics: List[Metric]) -> Dict[str, List[float]]:
+        metric_scores = {}
+        for metric in split_metrics:
+            try:
+                metric_scores[metric.tag] = [metric(inference_batch)]
+            except Exception as e:
+                raise MetricCalculationError(f"Error during calculation of metric {metric.tag}") from e
+        return metric_scores
+
+    def _calculate_loss_scores(self, forward_batch: InferenceResultBatch, split_loss_funs: Dict[str, Loss]) -> Dict[str, List[float]]:
+        loss_scores = {}
+        for loss_key, loss_fun in split_loss_funs.items():
+            try:
+                loss_scores[loss_key] = self._get_batch_loss(loss_fun, forward_batch)
+            except Exception as e:
+                raise LossCalculationError("Error during calculation of loss {loss_key}") from e
+
+        return loss_scores
+
+    def _get_batch_loss(self, loss_fun: Loss, forward_batch: InferenceResultBatch) -> List[torch.Tensor]:
+        loss = loss_fun(forward_batch)
+        loss = [loss.sum()]
+        return loss
diff --git a/src/ml_gym/gym/evaluator.py b/src/ml_gym/gym/evaluators/evaluator.py
similarity index 86%
rename from src/ml_gym/gym/evaluator.py
rename to src/ml_gym/gym/evaluators/evaluator.py
index 2cf29349..1e639745 100644
--- a/src/ml_gym/gym/evaluator.py
+++ b/src/ml_gym/gym/evaluators/evaluator.py
@@ -12,7 +12,6 @@
 from ml_gym.models.nn.net import NNModel
 from ml_gym.loss_functions.loss_functions import Loss
 import tqdm
-from ml_gym.util.logger import ConsoleLogger
 import numpy as np
 from ml_gym.gym.predict_postprocessing_component import PredictPostprocessingComponent
 from ml_gym.error_handling.exception import BatchStateError, EvaluationError, MetricCalculationError, LossCalculationError
@@ -27,17 +26,11 @@ def evaluate(self, model: NNModel, device: torch.device, epoch: int) -> List[Eva
 class Evaluator(AbstractEvaluator):
     def __init__(self, eval_component: 'EvalComponentIF'):
         self.eval_component = eval_component
-        self.current_epoch = -1
-        self.num_epochs = -1
 
-    def set_num_epochs(self, num_epochs: int):
-        self.num_epochs = num_epochs
-
-    def evaluate(self, model: NNModel, device: torch.device, current_epoch: int, num_epochs: int,
-                 epoch_result_callback_fun: Callable = None,
+    def evaluate(self, model: NNModel, device: torch.device, epoch_result_callback_fun: Callable = None,
                  batch_processed_callback_fun: Callable = None) -> List[EvaluationBatchResult]:
-        self.current_epoch = current_epoch
-        self.num_epochs = num_epochs
+        model = model.eval()
+
         # returns a EvaluationBatchResult for each split
         evaluation_batch_results = self.eval_component.evaluate(model, device, epoch_result_callback_fun=epoch_result_callback_fun,
                                                                 batch_processed_callback_fun=batch_processed_callback_fun)
@@ -56,7 +49,7 @@ class EvalComponent(EvalComponentIF):
     """This thing always comes with batteries included, i.e., datasets, loss functions etc. are all already stored in here."""
 
     def __init__(self, inference_component: InferenceComponent, post_processors: Dict[str, PredictPostprocessingComponent], metrics: List[Metric],
-                 loss_funs: Dict[str, Loss], dataset_loaders: Dict[str, DatasetLoader], train_split_name: str, show_progress: bool = False,
+                 loss_funs: Dict[str, Loss], dataset_loaders: Dict[str, DatasetLoader], show_progress: bool = False,
                  cpu_target_subscription_keys: List[str] = None, cpu_prediction_subscription_keys: List[Union[str, List]] = None,
                  metrics_computation_config: List[Dict] = None, loss_computation_config: List[Dict] = None):
         self.loss_funs = loss_funs
@@ -65,11 +58,9 @@ def __init__(self, inference_component: InferenceComponent, post_processors: Dic
         self.post_processors = post_processors
         self.metrics = metrics
         self.dataset_loaders = dataset_loaders
-        self.train_split_name = train_split_name
         self.show_progress = show_progress
         self.cpu_target_subscription_keys = cpu_target_subscription_keys
         self.cpu_prediction_subscription_keys = cpu_prediction_subscription_keys
-        self.logger = ConsoleLogger("logger_eval_component")
         # determines which metrics are applied to which splits (metric_key to split list)
         self.metrics_computation_config = None if metrics_computation_config is None else {
             m["metric_tag"]: m["applicable_splits"] for m in metrics_computation_config}
@@ -86,8 +77,7 @@ def evaluate_dataset_split(self, model: NNModel, device: torch.device, split_nam
                                dataset_loader: DatasetLoader, epoch_result_callback_fun: Callable = None,
                                batch_processed_callback_fun: Callable = None) -> EvaluationBatchResult:
         dataset_loader.device = device
-        dataset_loader_iterator = tqdm.tqdm(
-            dataset_loader, desc=f"Evaluating {dataset_loader.dataset_name} - {split_name}") if self.show_progress else dataset_loader
+        dataset_loader_iterator = tqdm.tqdm(dataset_loader, desc=f"Evaluating {dataset_loader.dataset_name} - {split_name}") if self.show_progress else dataset_loader
         post_processors = self.post_processors[split_name] + self.post_processors["default"]
 
         # calc losses
@@ -102,7 +92,6 @@ def evaluate_dataset_split(self, model: NNModel, device: torch.device, split_nam
         inference_result_batches_cpu = []
         num_batches = len(dataset_loader_iterator)
         processed_batches = 0
-        update_lag = max(1, int(num_batches/10))
         for batch in dataset_loader_iterator:
             inference_result_batch = self.forward_batch(dataset_batch=batch, model=model, device=device, postprocessors=post_processors)
             batch_loss = self._calculate_loss_scores(inference_result_batch, split_loss_funs)
@@ -112,13 +101,12 @@ def evaluate_dataset_split(self, model: NNModel, device: torch.device, split_nam
                                                                 device=torch.device("cpu"))
             inference_result_batches_cpu.append(irb_filtered)
             processed_batches += 1
-            if batch_processed_callback_fun is not None and (processed_batches % update_lag == 0 or processed_batches == num_batches):
-                splits = [d.dataset_tag for _, d in self.dataset_loaders.items()]
-                batch_processed_callback_fun(status="evaluation",
-                                             num_batches=num_batches,
-                                             current_batch=processed_batches,
-                                             splits=splits,
-                                             current_split=dataset_loader.dataset_tag)
+            splits = [d.dataset_tag for _, d in self.dataset_loaders.items()]
+            batch_processed_callback_fun(status="evaluation",
+                                         num_batches=num_batches,
+                                         current_batch=processed_batches,
+                                         splits=splits,
+                                         current_split=dataset_loader.dataset_tag)
 
         # calc metrics
         try:
@@ -153,7 +141,7 @@ def _get_metric_fun(self, identifier: str, target_subscription: Enum, prediction
 
     def forward_batch(self, dataset_batch: DatasetBatch, model: NNModel, device: torch.device, postprocessors: List[PredictPostProcessingIF]) -> InferenceResultBatch:
         model = model.to(device)
-        dataset_batch.to_device(device)
+        dataset_batch.to(device)
         inference_result_batch = self.inference_component.predict(model, dataset_batch, postprocessors)
         return inference_result_batch
 
diff --git a/src/ml_gym/gym/gym.py b/src/ml_gym/gym/gym.py
index 8a1728da..58d81abf 100644
--- a/src/ml_gym/gym/gym.py
+++ b/src/ml_gym/gym/gym.py
@@ -1,63 +1,142 @@
+from functools import partial
+from ml_gym.error_handling.exception import GymError
+from ml_gym.gym.gym_jobs.gym_job_factory import GymJobFactory
+from ml_gym.persistency.io import GridSearchAPIClientConstructable
 from ml_gym.persistency.logging import JobStatusLogger, MLgymStatusLoggerCollectionConstructable
 import torch
-from typing import List
+from typing import Callable, List
 from ml_gym.multiprocessing.pool import Pool, Job
 from ml_gym.blueprints.blue_prints import BluePrint
-from ml_gym.gym.jobs import AbstractGymJob
 from ml_gym.util.devices import get_devices
 import tqdm
-
+from abc import ABC, abstractmethod
+from enum import Enum
+from accelerate import Accelerator
 
 torch.multiprocessing.set_start_method('spawn', force=True)
 
 
-class Gym:
-    def __init__(self, job_id_prefix: str, logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
-                 process_count: int = 1, device_ids: List[int] = None, log_std_to_file: bool = True):
-        self.devices = get_devices(device_ids)
+class Gym(ABC):
+    def __init__(self, logger_collection_constructable: MLgymStatusLoggerCollectionConstructable):
         self.job_status_logger = JobStatusLogger(logger_collection_constructable.construct())
-        self.log_std_to_file = log_std_to_file
-        self.pool = Pool(num_processes=process_count, devices=self.devices, logger_collection_constructable=logger_collection_constructable)
-        self.jobs: List[Job] = []
-        self.job_id_prefix = job_id_prefix
-        self.job_counter = 0
 
-    def run(self, parallel=True):
+    @abstractmethod
+    def run(self, blueprints: List[BluePrint]):
         """Executes the jobs. Note that this function blocks until all jobs have been executed!
-
-        Args:
-            parallel (bool, optional): When set to True, jobs are run in parallel in the multiprocessing environment. Defaults to True.
         """
-        if parallel:
-            for _ in range(len(self.jobs)):
-                job = self.jobs.pop(0)
-                self.pool.add_job(job)
+        raise NotImplementedError
 
-            self.pool.run()
-        else:
-            for _ in tqdm.tqdm(range(len(self.jobs)), desc="Models trained"):
-                job = self.jobs.pop(0)
-                self.work(job, self.devices[0])
-
-    def add_blueprint(self, blueprint: BluePrint) -> int:
-        job = Job(job_id=f"{blueprint.grid_search_id}-{self.job_counter}", fun=Gym._run_job, blueprint=blueprint, param_dict={"log_std_to_file": self.log_std_to_file})
-        self.job_counter += 1
-        self.jobs.append(job)
-        return job.job_id
-
-    def add_blueprints(self, blueprints: List[BluePrint]):
-        for blueprint in blueprints:
-            job_id = self.add_blueprint(blueprint)
+    def get_jobs_from_blueprints(self, blueprints: List[BluePrint], exec_fun: Callable) -> List[Job]:
+        jobs = []
+        for job_id, blueprint in enumerate(blueprints):
+            job = Job(job_id=f"{blueprint.grid_search_id}-{job_id}", fun=exec_fun, blueprint=blueprint)
+            jobs.append(job)
             self.job_status_logger.log_experiment_config(grid_search_id=blueprint.grid_search_id,
                                                          experiment_id=blueprint.experiment_id,
                                                          job_id=job_id,
                                                          config=blueprint.config)
+        return jobs
+
+
+class ParallelSingleNodeGym(Gym):
+    def __init__(self, logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
+                 exec_fun: Callable, process_count: int = 1, device_ids: List[int] = None):
+        super().__init__(logger_collection_constructable)
+        self.devices = get_devices(device_ids)
+        self.pool = Pool(num_processes=process_count, devices=self.devices, logger_collection_constructable=logger_collection_constructable)
+        self.exec_fun = exec_fun
+
+    def run(self, blueprints: List[BluePrint]):
+        """Executes the jobs. Note that this function blocks until all jobs have been executed!
+        """
+        jobs = self.get_jobs_from_blueprints(blueprints, exec_fun=self.exec_fun)
+        for job in jobs:
+            self.pool.add_job(job)
+        self.pool.run()
+
+
+class SequentialGym(Gym):
+    def __init__(self, logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
+                 exec_fun: Callable, device_id: int = None):
+        super().__init__(logger_collection_constructable)
+        self.exec_fun = exec_fun
+        self.device = None if device_id is None else torch.device(device_id)
+
+    def run(self, blueprints: List[BluePrint]):
+        """Executes the jobs. Note that this function blocks until all jobs have been executed!
+        """
+
+        jobs = self.get_jobs_from_blueprints(blueprints, exec_fun=self.exec_fun)
+        for job in tqdm.tqdm(jobs, desc="Models trained"):
+            job.execute()
+        self.job_status_logger.disconnect()
+
+    def work(self, job: Job):
+        # we don't set the device in job because every job executed on the main thread should be handled via accelerate
+        job.execute(device=self.device)
+
+
+class GymType(Enum):
+    SEQUENTIAL_GYM = SequentialGym
+    PARALLEL_SINGLE_NODE_GYM = ParallelSingleNodeGym
+
+
+class GymFactory:
+
+    @staticmethod
+    def get_sequential_gym(logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
+                           gs_restful_api_client_constructable: GridSearchAPIClientConstructable,
+                           num_epochs: int, device_id: int = None,
+                           num_batches_per_epoch: int = None,
+                           accelerator: Accelerator = None) -> Gym:
+        if accelerator is not None:
+            exec_fun = partial(GymFactory._run_accelerate_gym_job, logger_collection_constructable=logger_collection_constructable,
+                               gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                               num_epochs=num_epochs, num_batches_per_epoch=num_batches_per_epoch, accelerator=accelerator)
+            if device_id is not None:
+                raise GymError("You cannot run the accelerate env on a specific device id")
+        else:
+            exec_fun = partial(GymFactory._run_standard_gym_job, logger_collection_constructable=logger_collection_constructable,
+                               gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                               num_epochs=num_epochs, num_batches_per_epoch=num_batches_per_epoch)
+
+        return SequentialGym(logger_collection_constructable=logger_collection_constructable,
+                             exec_fun=exec_fun,
+                             device_id=device_id)
+
+    @staticmethod
+    def get_parallel_single_node_gym(logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
+                                     gs_restful_api_client_constructable: GridSearchAPIClientConstructable,
+                                     num_epochs: int, process_count: int = 1, device_ids: List[int] = None,
+                                     num_batches_per_epoch: int = None) -> Gym:
+
+        exec_fun = partial(GymFactory._run_standard_gym_job, logger_collection_constructable=logger_collection_constructable,
+                           gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                           num_epochs=num_epochs, num_batches_per_epoch=num_batches_per_epoch)
+
+        return ParallelSingleNodeGym(logger_collection_constructable=logger_collection_constructable,
+                                     process_count=process_count, device_ids=device_ids, exec_fun=exec_fun)
 
     @staticmethod
-    def _run_job(blueprint: BluePrint, device: torch.device, log_std_to_file: bool) -> AbstractGymJob:
-        gym_job = AbstractGymJob.from_blue_print(blueprint, device=device)
+    def _run_standard_gym_job(blueprint: BluePrint, device: torch.device, num_epochs: int,
+                              logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
+                              gs_restful_api_client_constructable: GridSearchAPIClientConstructable,
+                              num_batches_per_epoch: int = None):
+        gym_job = GymJobFactory.get_gym_job_from_blueprint(blueprint, device=device, num_epochs=num_epochs,
+                                                           logger_collection_constructable=logger_collection_constructable,
+                                                           gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                                                           num_batches_per_epoch=num_batches_per_epoch)
         return gym_job.execute(device=device)
 
-    def work(self, job: Job, device: torch.device):
-        job.device = device
-        job.execute()
+    @staticmethod
+    def _run_accelerate_gym_job(blueprint: BluePrint, num_epochs: int,
+                                logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
+                                gs_restful_api_client_constructable: GridSearchAPIClientConstructable,
+                                accelerator: Accelerator, num_batches_per_epoch: int = None):
+        gym_job = GymJobFactory.get_accelerate_gymjob_from_blueprint(blueprint,
+                                                                     num_epochs=num_epochs,
+                                                                     logger_collection_constructable=logger_collection_constructable,
+                                                                     gs_restful_api_client_constructable=gs_restful_api_client_constructable,
+                                                                     num_batches_per_epoch=num_batches_per_epoch,
+                                                                     accelerator=accelerator)
+        return gym_job.execute()
diff --git a/src/ml_gym/gym/gym_jobs/__init__.py b/src/ml_gym/gym/gym_jobs/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/ml_gym/gym/gym_jobs/accelerate_gym_job.py b/src/ml_gym/gym/gym_jobs/accelerate_gym_job.py
new file mode 100644
index 00000000..c1c628df
--- /dev/null
+++ b/src/ml_gym/gym/gym_jobs/accelerate_gym_job.py
@@ -0,0 +1,112 @@
+from ml_gym.data_handling.dataset_loader import DatasetLoaderFactory
+from ml_gym.early_stopping.early_stopping_strategies import EarlyStoppingIF
+from ml_gym.gym.gym_jobs.gym_job import AbstractGymJob
+from ml_gym.gym.trainers.accelerate_trainer import AccelerateTrainer
+from ml_gym.models.nn.net import NNModel
+from ml_gym.gym.evaluators.evaluator import Evaluator
+from ml_gym.modes import RunMode
+from ml_gym.optimizers.lr_schedulers import LRSchedulerAdapter
+from ml_gym.optimizers.optimizer import OptimizerAdapter
+from typing import List
+from ml_gym.batching.batch import EvaluationBatchResult
+from ml_gym.persistency.logging import ExperimentStatusLogger
+from functools import partial
+from ml_gym.persistency.io import GridSearchAPIClientIF, CheckpointResource
+import pickle
+from ml_gym.checkpointing.checkpointing import CheckpointingIF
+from accelerate import Accelerator
+
+
+class AccelerateGymJob(AbstractGymJob):
+
+    def __init__(self, experiment_status_logger: ExperimentStatusLogger, gs_api_client: GridSearchAPIClientIF,
+                 grid_search_id: str, experiment_id: int, run_mode: RunMode, num_epochs: int,
+                 model: NNModel, optimizer: OptimizerAdapter, trainer: AccelerateTrainer, evaluator: Evaluator,
+                 checkpointing_strategy: CheckpointingIF, accelerator: Accelerator, early_stopping_strategy: EarlyStoppingIF = None,
+                 warm_start_epoch: int = 0, lr_scheduler: LRSchedulerAdapter = None, num_batches_per_epoch: int = None):
+        super().__init__(experiment_status_logger=experiment_status_logger, gs_api_client=gs_api_client, grid_search_id=grid_search_id,
+                         experiment_id=experiment_id, run_mode=run_mode, num_epochs=num_epochs, model=model, optimizer=optimizer,
+                         trainer=trainer, evaluator=evaluator, checkpointing_strategy=checkpointing_strategy,
+                         early_stopping_strategy=early_stopping_strategy, warm_start_epoch=warm_start_epoch, lr_scheduler=lr_scheduler,
+                         num_batches_per_epoch=num_batches_per_epoch)
+        self.num_batches_per_epoch = num_batches_per_epoch
+        self.accelerator = accelerator
+
+    def execute(self):
+        """ Executes the job
+
+        """
+        self._execution_method()
+        self._experiment_status_logger.disconnect()
+
+    def _evaluation_step(self, current_epoch: int) -> List[EvaluationBatchResult]:
+        partial_batch_processed_callback = partial(self.batch_processed_callback, num_epochs=self.num_epochs,
+                                                   current_epoch=current_epoch,
+                                                   experiment_status_logger=self._experiment_status_logger)
+        partial_epoch_result_callback = partial(self.epoch_result_callback, current_epoch=current_epoch,
+                                                experiment_status_logger=self._experiment_status_logger)
+
+        evaluation_results = self.evaluator.evaluate(model=self.model,
+                                                     accelerator=self.accelerator,
+                                                     batch_processed_callback_fun=partial_batch_processed_callback,
+                                                     epoch_result_callback_fun=partial_epoch_result_callback)
+
+        return evaluation_results
+
+    def _execute_train(self):
+        self.optimizer.register_model_params(model_params=dict(self.model.named_parameters()))
+        self.lr_scheduler.register_optimizer(optimizer=self.optimizer)
+
+        self.model, self.optimizer, self.trainer, self.evaluator, self.lr_scheduler, train_loader = self.accelerator.prepare(self.model,
+                                                                                                                             self.optimizer,
+                                                                                                                             self.trainer,
+                                                                                                                             self.evaluator,
+                                                                                                                             self.lr_scheduler,
+                                                                                                                             self.trainer.train_loader)
+        self.trainer.train_loader = DatasetLoaderFactory.get_data_loader_shard_wrapper(data_loader_shard=train_loader,
+                                                                                       dataset_name=self.trainer.train_loader.dataset_name,
+                                                                                       dataset_tag=self.trainer.train_loader.dataset_tag)
+
+        self.evaluator.eval_component.dataset_loaders = {key: DatasetLoaderFactory.get_data_loader_shard_wrapper(
+            data_loader_shard=self.accelerator.prepare(data_loader),
+            dataset_name=data_loader.dataset_name,
+            dataset_tag=data_loader.dataset_tag) for key, data_loader in self.evaluator.eval_component.dataset_loaders.items()}
+
+        partial_batch_done_callback = partial(self.batch_processed_callback, experiment_status_logger=self._experiment_status_logger)
+        def evaluation_step_routine(current_epoch: int): return self._evaluation_step(current_epoch=current_epoch)
+        partial_train_epoch_done_callback = partial(self.train_epoch_done_callback, evaluation_step_routine=evaluation_step_routine, 
+                                                    accelerator=self.accelerator)
+
+        model = self.trainer.train(num_epochs=self.num_epochs, model=self.model, optimizer=self.optimizer,
+                                   batch_done_callback_fun=partial_batch_done_callback,
+                                   epoch_done_callback=partial_train_epoch_done_callback, accelerator=self.accelerator,
+                                   num_batches_per_epoch=self.num_batches_per_epoch)
+        self.accelerator.free_memory()
+
+    def _execute_warm_start(self):
+        if self.current_epoch > 0:
+            model_state = pickle.loads(self.gs_api_client.get_checkpoint_resource(grid_search_id=self.grid_search_id,
+                                                                                  experiment_id=self.experiment_id,
+                                                                                  checkpoint_id=self.current_epoch,
+                                                                                  checkpoint_resource=CheckpointResource.model))
+            self.model.load_state_dict(model_state)
+
+            optimizer_state = pickle.loads(self.gs_api_client.get_checkpoint_resource(grid_search_id=self.grid_search_id,
+                                                                                      experiment_id=self.experiment_id,
+                                                                                      checkpoint_id=self.current_epoch,
+                                                                                      checkpoint_resource=CheckpointResource.optimizer))
+            self.optimizer.load_state_dict(optimizer_state)
+
+            lr_scheduler_state = pickle.loads(self.gs_api_client.get_checkpoint_resource(grid_search_id=self.grid_search_id,
+                                                                                         experiment_id=self.experiment_id,
+                                                                                         checkpoint_id=self.current_epoch,
+                                                                                         checkpoint_resource=CheckpointResource.lr_scheduler))
+            self.lr_scheduler.load_state_dict(lr_scheduler_state)
+
+            stateful_component_state = pickle.loads(self.gs_api_client.get_checkpoint_resource(grid_search_id=self.grid_search_id,
+                                                                                               experiment_id=self.experiment_id,
+                                                                                               checkpoint_id=self.current_epoch,
+                                                                                               checkpoint_resource=CheckpointResource.stateful_components))
+            self.set_state(stateful_component_state)
+
+        self._execute_train()
diff --git a/src/ml_gym/gym/gym_jobs/gym_job.py b/src/ml_gym/gym/gym_jobs/gym_job.py
new file mode 100644
index 00000000..6da8e1b2
--- /dev/null
+++ b/src/ml_gym/gym/gym_jobs/gym_job.py
@@ -0,0 +1,119 @@
+from abc import abstractmethod
+from typing import Callable, List
+from ml_gym.batching.batch import EvaluationBatchResult
+from ml_gym.checkpointing.checkpointing import CheckpointingIF, CheckpointingInstruction
+from ml_gym.early_stopping.early_stopping_strategies import EarlyStoppingIF
+from ml_gym.error_handling.exception import EarlyStoppingCriterionFulfilledError
+from ml_gym.gym.evaluators.evaluator import Evaluator
+from ml_gym.gym.stateful_components import StatefulComponent
+from ml_gym.gym.trainers.standard_trainer import Trainer
+from ml_gym.models.nn.net import NNModel
+from ml_gym.modes import RunMode
+from ml_gym.optimizers.lr_scheduler_factory import LRSchedulerFactory
+from ml_gym.optimizers.lr_schedulers import LRSchedulerAdapter
+from ml_gym.optimizers.optimizer import OptimizerAdapter
+from ml_gym.persistency.io import GridSearchAPIClientIF
+from ml_gym.persistency.logging import ExperimentStatusLogger
+import torch
+from accelerate import Accelerator
+
+
+class AbstractGymJob(StatefulComponent):
+
+    def __init__(self, experiment_status_logger: ExperimentStatusLogger, gs_api_client: GridSearchAPIClientIF,
+                 grid_search_id: str, experiment_id: int, run_mode: RunMode, num_epochs: int,
+                 model: NNModel, optimizer: OptimizerAdapter, trainer: Trainer, evaluator: Evaluator,
+                 checkpointing_strategy: CheckpointingIF, early_stopping_strategy: EarlyStoppingIF = None,
+                 warm_start_epoch: int = 0, lr_scheduler: LRSchedulerAdapter = None, num_batches_per_epoch: int = None):
+        super().__init__()
+        # logging
+        self._experiment_status_logger = experiment_status_logger
+        self.gs_api_client = gs_api_client
+
+        # identifiers
+        self.grid_search_id = grid_search_id
+        self.experiment_id = experiment_id
+
+        # configuration
+        self.run_mode = run_mode
+        self.num_epochs = num_epochs
+        self.num_batches_per_epoch = num_batches_per_epoch
+        self.warm_start_epoch = warm_start_epoch
+        if run_mode == RunMode.TRAIN:
+            self._execution_method = self._execute_train
+        elif run_mode == RunMode.WARM_START:
+            self._execution_method = self._execute_warm_start
+        else:
+            raise NotImplementedError
+
+        # components
+        self.model = model
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler if lr_scheduler is not None else LRSchedulerFactory.get_lr_scheduler("dummy")
+        self.evaluator = evaluator
+        self.trainer = trainer
+
+        self.checkpointing_strategy = checkpointing_strategy
+        self.early_stopping_strategy = early_stopping_strategy
+
+    @abstractmethod
+    def execute(self, device: torch.device):
+        raise NotImplementedError
+
+    def run_checkpointing(self, checkpoint_instruction: CheckpointingInstruction, current_epoch: int):
+        """
+        Create and delete checkpoints for each epoch in experiments.
+
+        :params:
+            checkpoint_instruction: CheckpointingInstruction object
+        """
+
+        if checkpoint_instruction.save_current:
+            payload_dict = {
+                "epoch": current_epoch,
+                "model": self.model.state_dict(),
+                "optimizer": self.optimizer.state_dict(),
+                "lr_scheduler": self.lr_scheduler.state_dict(),
+                "stateful_components": self.get_state(),
+            }
+
+            self.gs_api_client.add_checkpoint_resource(
+                grid_search_id=self.grid_search_id, experiment_id=self.experiment_id, payload=payload_dict
+            )
+
+        for epoch in checkpoint_instruction.checkpoints_to_delete:
+            self.gs_api_client.delete_checkpoint_resource(grid_search_id=self.grid_search_id, experiment_id=self.experiment_id, epoch=epoch)
+
+    @staticmethod
+    def batch_processed_callback(status: str, experiment_status_logger: ExperimentStatusLogger, num_batches: int,
+                                 current_batch: int, splits: List[str], current_split: str, num_epochs: int, current_epoch: int):
+        if (current_batch % max(1, int(num_batches/10))) == 0:  # TODO make update period configurable
+            experiment_status_logger.log_experiment_status(status=status,
+                                                           num_epochs=num_epochs,
+                                                           current_epoch=current_epoch,
+                                                           splits=splits,
+                                                           current_split=current_split,
+                                                           num_batches=num_batches,
+                                                           current_batch=current_batch)
+
+    @staticmethod
+    def epoch_result_callback(experiment_status_logger: ExperimentStatusLogger, evaluation_result: EvaluationBatchResult,
+                              current_epoch: int):
+        experiment_status_logger.log_evaluation_results(evaluation_result, current_epoch)
+
+    def train_epoch_done_callback(self, num_epochs: int, current_epoch: int, model: NNModel, evaluation_step_routine: Callable,
+                                  accelerator: Accelerator = None):
+        evaluation_results = evaluation_step_routine(current_epoch=current_epoch)
+        if current_epoch > 0:
+            self.lr_scheduler.step()
+        if accelerator is not None and accelerator.is_main_process:
+            checkpointing_instruction = self.checkpointing_strategy.get_model_checkpoint_instruction(num_epochs=num_epochs,
+                                                                                                     current_epoch=current_epoch,
+                                                                                                     evaluation_result=evaluation_results)
+            self.run_checkpointing(checkpointing_instruction, current_epoch=current_epoch)
+        if self.early_stopping_strategy.is_stopping_criterion_fulfilled(current_epoch=current_epoch,
+                                                                        evaluation_results=evaluation_results):
+            raise EarlyStoppingCriterionFulfilledError
+
+    def _evaluation_step(self, device: torch.device) -> List[EvaluationBatchResult]:
+        raise NotImplementedError
diff --git a/src/ml_gym/gym/gym_jobs/gym_job_factory.py b/src/ml_gym/gym/gym_jobs/gym_job_factory.py
new file mode 100644
index 00000000..94034f78
--- /dev/null
+++ b/src/ml_gym/gym/gym_jobs/gym_job_factory.py
@@ -0,0 +1,59 @@
+from ml_gym.blueprints.blue_prints import BluePrint
+from ml_gym.gym.gym_jobs.accelerate_gym_job import AccelerateGymJob
+from ml_gym.gym.gym_jobs.gym_job import AbstractGymJob
+from ml_gym.gym.gym_jobs.standard_gym_job import StandardGymJob
+from ml_gym.persistency.io import GridSearchAPIClientConstructable
+from ml_gym.persistency.logging import ExperimentStatusLogger, MLgymStatusLoggerCollectionConstructable
+import torch
+from accelerate import Accelerator
+
+
+class GymJobFactory:
+
+    @staticmethod
+    def get_gym_job_from_blueprint(blueprint: BluePrint, device: torch.device, num_epochs: int,
+                                   logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
+                                   gs_restful_api_client_constructable: GridSearchAPIClientConstructable,
+                                   num_batches_per_epoch: int = None) -> AbstractGymJob:
+        components = blueprint.construct(device)
+
+        logger_collection = logger_collection_constructable.construct()
+        experiment_status_logger = ExperimentStatusLogger(logger=logger_collection, grid_search_id=blueprint.grid_search_id,
+                                                          experiment_id=blueprint.experiment_id)
+        gs_api_client = gs_restful_api_client_constructable.construct()
+
+        gym_job = StandardGymJob(run_mode=blueprint.run_mode,
+                                 grid_search_id=blueprint.grid_search_id,
+                                 experiment_id=blueprint.experiment_id,
+                                 num_epochs=num_epochs,
+                                 warm_start_epoch=blueprint.warm_start_epoch,
+                                 experiment_status_logger=experiment_status_logger,
+                                 gs_api_client=gs_api_client,
+                                 num_batches_per_epoch=num_batches_per_epoch,
+                                 **components)
+        return gym_job
+
+    @staticmethod
+    def get_accelerate_gymjob_from_blueprint(blueprint: BluePrint, num_epochs: int,
+                                             logger_collection_constructable: MLgymStatusLoggerCollectionConstructable,
+                                             gs_restful_api_client_constructable: GridSearchAPIClientConstructable,
+                                             accelerator: Accelerator,
+                                             num_batches_per_epoch: int = None) -> AbstractGymJob:
+        components = blueprint.construct()
+
+        logger_collection = logger_collection_constructable.construct()
+        experiment_status_logger = ExperimentStatusLogger(logger=logger_collection, grid_search_id=blueprint.grid_search_id,
+                                                          experiment_id=blueprint.experiment_id)
+        gs_api_client = gs_restful_api_client_constructable.construct()
+
+        gym_job = AccelerateGymJob(run_mode=blueprint.run_mode,
+                                   grid_search_id=blueprint.grid_search_id,
+                                   experiment_id=blueprint.experiment_id,
+                                   num_epochs=num_epochs,
+                                   warm_start_epoch=blueprint.warm_start_epoch,
+                                   experiment_status_logger=experiment_status_logger,
+                                   gs_api_client=gs_api_client,
+                                   num_batches_per_epoch=num_batches_per_epoch,
+                                   accelerator=accelerator,
+                                   **components)
+        return gym_job
diff --git a/src/ml_gym/gym/gym_jobs/standard_gym_job.py b/src/ml_gym/gym/gym_jobs/standard_gym_job.py
new file mode 100644
index 00000000..86f0fbad
--- /dev/null
+++ b/src/ml_gym/gym/gym_jobs/standard_gym_job.py
@@ -0,0 +1,104 @@
+from ml_gym.early_stopping.early_stopping_strategies import EarlyStoppingIF
+from ml_gym.gym.gym_jobs.gym_job import AbstractGymJob
+from ml_gym.models.nn.net import NNModel
+from ml_gym.gym.evaluators.evaluator import Evaluator
+from ml_gym.gym.trainers.standard_trainer import Trainer
+from ml_gym.modes import RunMode
+from ml_gym.optimizers.lr_schedulers import LRSchedulerAdapter
+from ml_gym.optimizers.optimizer import OptimizerAdapter
+import torch
+from typing import List
+from ml_gym.batching.batch import EvaluationBatchResult
+from ml_gym.persistency.logging import ExperimentStatusLogger
+from functools import partial
+from ml_gym.persistency.io import GridSearchAPIClientIF, CheckpointResource
+import pickle
+from ml_gym.checkpointing.checkpointing import CheckpointingIF
+from ml_gym.error_handling.exception import EarlyStoppingCriterionFulfilledError
+
+
+class StandardGymJob(AbstractGymJob):
+
+    def __init__(self, experiment_status_logger: ExperimentStatusLogger, gs_api_client: GridSearchAPIClientIF,
+                 grid_search_id: str, experiment_id: int, run_mode: RunMode, num_epochs: int,
+                 model: NNModel, optimizer: OptimizerAdapter, trainer: Trainer, evaluator: Evaluator,
+                 checkpointing_strategy: CheckpointingIF, early_stopping_strategy: EarlyStoppingIF = None,
+                 warm_start_epoch: int = 0, lr_scheduler: LRSchedulerAdapter = None, num_batches_per_epoch: int = None):
+        super().__init__(experiment_status_logger=experiment_status_logger, gs_api_client=gs_api_client, grid_search_id=grid_search_id,
+                         experiment_id=experiment_id, run_mode=run_mode, num_epochs=num_epochs, model=model, optimizer=optimizer,
+                         trainer=trainer, evaluator=evaluator, checkpointing_strategy=checkpointing_strategy,
+                         early_stopping_strategy=early_stopping_strategy, warm_start_epoch=warm_start_epoch, lr_scheduler=lr_scheduler,
+                         num_batches_per_epoch=num_batches_per_epoch)
+
+    def execute(self, device: torch.device):
+        """ Executes the job
+
+        Args:
+            device: torch device either CPUs or a specified GPU
+        """
+        self._execution_method(device)
+        self._experiment_status_logger.disconnect()
+
+    def _train_step(self, device: torch.device) -> NNModel:
+        partial_batch_processed_callback = partial(AbstractGymJob.batch_processed_callback, num_epochs=self.num_epochs,
+                                                   experiment_status_logger=self._experiment_status_logger)
+        model = self.trainer.train_epoch(self.model, self.optimizer, device,
+                                         batch_processed_callback_fun=partial_batch_processed_callback)
+        return model
+
+    def _evaluation_step(self, current_epoch: int, device: torch.device) -> List[EvaluationBatchResult]:
+        partial_batch_processed_callback = partial(AbstractGymJob.batch_processed_callback, num_epochs=self.num_epochs,
+                                                   current_epoch=current_epoch,
+                                                   experiment_status_logger=self._experiment_status_logger)
+        partial_epoch_result_callback = partial(self.epoch_result_callback, experiment_status_logger=self._experiment_status_logger,
+                                                current_epoch=current_epoch)
+
+        evaluation_results = self.evaluator.evaluate(model=self.model,
+                                                     device=device,
+                                                     batch_processed_callback_fun=partial_batch_processed_callback,
+                                                     epoch_result_callback_fun=partial_epoch_result_callback)
+
+        return evaluation_results
+
+    def _execute_train(self, device: torch.device):
+        self.optimizer.register_model_params(model_params=dict(self.model.named_parameters()))
+        self.lr_scheduler.register_optimizer(optimizer=self.optimizer)
+
+        partial_batch_done_callback = partial(self.batch_processed_callback, experiment_status_logger=self._experiment_status_logger)
+        def evaluation_step_routine(current_epoch: int): return self._evaluation_step(device=device, current_epoch=current_epoch)
+        partial_train_epoch_done_callback = partial(self.train_epoch_done_callback, evaluation_step_routine=evaluation_step_routine)
+
+        try:
+            model = self.trainer.train(num_epochs=self.num_epochs, model=self.model, optimizer=self.optimizer, device=device,
+                                       batch_done_callback_fun=partial_batch_done_callback,
+                                       epoch_done_callback=partial_train_epoch_done_callback,
+                                       num_batches_per_epoch=self.num_batches_per_epoch)
+        except EarlyStoppingCriterionFulfilledError:
+            print(f"Early stopping criterion matched. Stopping training.")
+
+    def _execute_warm_start(self, device: torch.device, warm_start_epoch: int):
+        model_state = pickle.loads(self.gs_api_client.get_checkpoint_resource(grid_search_id=self.grid_search_id,
+                                                                              experiment_id=self.experiment_id,
+                                                                              checkpoint_id=warm_start_epoch,
+                                                                              checkpoint_resource=CheckpointResource.model))
+        self.model.load_state_dict(model_state)
+
+        optimizer_state = pickle.loads(self.gs_api_client.get_checkpoint_resource(grid_search_id=self.grid_search_id,
+                                                                                  experiment_id=self.experiment_id,
+                                                                                  checkpoint_id=warm_start_epoch,
+                                                                                  checkpoint_resource=CheckpointResource.optimizer))
+        self.optimizer.load_state_dict(optimizer_state)
+
+        lr_scheduler_state = pickle.loads(self.gs_api_client.get_checkpoint_resource(grid_search_id=self.grid_search_id,
+                                                                                     experiment_id=self.experiment_id,
+                                                                                     checkpoint_id=warm_start_epoch,
+                                                                                     checkpoint_resource=CheckpointResource.lr_scheduler))
+        self.lr_scheduler.load_state_dict(lr_scheduler_state)
+
+        stateful_component_state = pickle.loads(self.gs_api_client.get_checkpoint_resource(grid_search_id=self.grid_search_id,
+                                                                                           experiment_id=self.experiment_id,
+                                                                                           checkpoint_id=warm_start_epoch,
+                                                                                           checkpoint_resource=CheckpointResource.stateful_components))
+        self.set_state(stateful_component_state)
+
+        self._execute_train(device)
diff --git a/src/ml_gym/gym/jobs.py b/src/ml_gym/gym/jobs.py
deleted file mode 100644
index 03c6347f..00000000
--- a/src/ml_gym/gym/jobs.py
+++ /dev/null
@@ -1,329 +0,0 @@
-from abc import abstractmethod
-from ml_gym.early_stopping.early_stopping_strategies import EarlyStoppingIF
-from ml_gym.models.nn.net import NNModel
-from ml_gym.gym.evaluator import Evaluator
-from ml_gym.gym.trainer import Trainer
-from ml_gym.modes import RunMode
-from ml_gym.optimizers.lr_scheduler_factory import LRSchedulerFactory
-from ml_gym.optimizers.lr_schedulers import LRSchedulerAdapter
-from ml_gym.optimizers.optimizer import OptimizerAdapter
-import torch
-from ml_gym.gym.stateful_components import StatefulComponent
-from typing import List, Dict, Any
-from ml_gym.util.logger import ConsoleLogger, LogLevel
-from ml_gym.batching.batch import EvaluationBatchResult
-from ml_gym.persistency.logging import ExperimentStatusLogger
-from functools import partial
-from ml_gym.persistency.io import GridSearchAPIClientIF, CheckpointResource
-import pickle
-from ml_gym.checkpointing.checkpointing import CheckpointingIF, CheckpointingInstruction
-
-
-class AbstractGymJob(StatefulComponent):
-    """
-    Gym Job Abstract Class
-
-    Declare Abstract and Static methods to be used in Gym Job Class.
-    """
-
-    def __init__(self, experiment_status_logger: ExperimentStatusLogger):
-        super().__init__()
-        self._experiment_status_logger = experiment_status_logger
-        # TODO: This is only a console logger that does not store on disk.
-        # Since GymJob is a StatefulComponent we would have to implement serialization
-        # for the queue in the QLogger. Solving this issue is left for the future
-        self.logger = ConsoleLogger("logger_gym_job")
-
-    @abstractmethod
-    def execute(self, device: torch.device):
-        raise NotImplementedError
-
-    @staticmethod
-    def from_blue_print(blue_print, device: torch.device) -> "AbstractGymJob":
-        return blue_print.construct(device)
-
-    # def restore_state_in_stateful_components(self, measurement_id: int):
-    #     state = DashifyReader.load_state(experiment_info=self.experiment_info, measurement_id=measurement_id)
-    #     self.set_state(state)
-
-    @staticmethod
-    def batch_processed_callback(
-        status: str,
-        experiment_status_logger: ExperimentStatusLogger,
-        num_batches: int,
-        current_batch: int,
-        splits: List[str],
-        current_split: str,
-        num_epochs: int,
-        current_epoch: int,
-    ):
-        experiment_status_logger.log_experiment_status(
-            status=status,
-            num_epochs=num_epochs,
-            current_epoch=current_epoch,
-            splits=splits,
-            current_split=current_split,
-            num_batches=num_batches,
-            current_batch=current_batch,
-        )
-
-    @staticmethod
-    def epoch_result_callback(
-        experiment_status_logger: ExperimentStatusLogger, evaluation_result: EvaluationBatchResult, current_epoch: int
-    ):
-        experiment_status_logger.log_evaluation_results(evaluation_result, current_epoch)
-
-
-class GymJob(AbstractGymJob):
-    """
-    GymJob Class
-
-    :params: AbstractGymJob object
-
-    The GymJob class is a CORE class in the mlGym which is responsible for running and logging jobs.
-    """
-
-    def __init__(
-        self,
-        grid_search_id: str,
-        experiment_id: int,
-        run_mode: RunMode,
-        model: NNModel,
-        optimizer: OptimizerAdapter,
-        trainer: Trainer,
-        evaluator: Evaluator,
-        num_epochs: int,
-        checkpointing_strategy: CheckpointingIF,
-        gs_api_client: GridSearchAPIClientIF,
-        experiment_status_logger: ExperimentStatusLogger = None,
-        early_stopping_strategy: EarlyStoppingIF = None,
-        warm_start_epoch: int = 0,
-        lr_scheduler: LRSchedulerAdapter = None,
-    ):
-        super().__init__(experiment_status_logger)
-        self.grid_search_id = grid_search_id
-        self.experiment_id = experiment_id
-        self.run_mode = run_mode
-        self.model = model
-        self.optimizer = optimizer
-        self.lr_scheduler = lr_scheduler if lr_scheduler is not None else LRSchedulerFactory.get_lr_scheduler("dummy")
-        # self.optimizer.register_model_params(dict(self.model.named_parameters()))
-        self.num_epochs = num_epochs
-        self.current_epoch = warm_start_epoch
-        self.evaluator = evaluator
-        self.trainer = trainer
-        self.checkpointing_strategy = checkpointing_strategy
-        self.early_stopping_strategy = early_stopping_strategy
-        self.gs_api_client = gs_api_client
-        if run_mode == RunMode.TRAIN:
-            self._execution_method = self._execute_train
-        elif run_mode == RunMode.WARM_START:
-            self._execution_method = self._execute_warm_start
-        else:
-            raise NotImplementedError
-
-    def _train_step(self, device: torch.device) -> NNModel:
-        # self.restore_state_in_stateful_components(epoch - 1)
-        # load model and optimizer
-        # model_state = DashifyReader.load_model_state(self._experiment_info, epoch - 1)
-        # self.model.load_state_dict(model_state)
-        # self.model.to(device)
-        # optimizer_state = DashifyReader.load_optimizer_state(self._experiment_info, epoch - 1)
-        # self.optimizer.load_state_dict(optimizer_state)
-        # train
-        partial_batch_processed_callback = partial(
-            self.batch_processed_callback,
-            num_epochs=self.num_epochs,
-            current_epoch=self.current_epoch,
-            experiment_status_logger=self._experiment_status_logger,
-        )
-        model = self.trainer.train_epoch(self.model, self.optimizer, device, batch_processed_callback_fun=partial_batch_processed_callback)
-        return model
-
-    def _evaluation_step(self, device: torch.device) -> List[EvaluationBatchResult]:
-        self.model.to(device)
-        partial_batch_processed_callback = partial(
-            self.batch_processed_callback,
-            num_epochs=self.num_epochs,
-            current_epoch=self.current_epoch,
-            experiment_status_logger=self._experiment_status_logger,
-        )
-        partial_epoch_result_callback = partial(
-            self.epoch_result_callback, current_epoch=self.current_epoch, experiment_status_logger=self._experiment_status_logger
-        )
-
-        evaluation_results = self.evaluator.evaluate(
-            model=self.model,
-            device=device,
-            current_epoch=self.current_epoch,
-            num_epochs=self.num_epochs,
-            batch_processed_callback_fun=partial_batch_processed_callback,
-            epoch_result_callback_fun=partial_epoch_result_callback,
-        )
-
-        return evaluation_results
-
-    def run_checkpointing(self, checkpoint_instruction: CheckpointingInstruction):
-        """
-        Create and Delete checkpoints for each epoch in experiments.
-
-        :params:
-            checkpoint_instruction: CheckpointingInstruction object
-        """
-
-        if checkpoint_instruction.save_current:
-            payload_dict = {
-                "epoch": self.current_epoch,
-                "model": self.model.state_dict(),
-                "optimizer": self.optimizer.state_dict(),
-                "lr_scheduler": self.lr_scheduler.state_dict(),
-                "stateful_components": self.get_state(),
-            }
-
-            self.gs_api_client.add_checkpoint_resource(
-                grid_search_id=self.grid_search_id, experiment_id=self.experiment_id, payload=payload_dict
-            )
-
-        for epoch in checkpoint_instruction.checkpoints_to_delete:
-            print(f"epoch to delete: {epoch}")
-            self.gs_api_client.delete_checkpoint_resource(grid_search_id=self.grid_search_id, experiment_id=self.experiment_id, epoch=epoch)
-
-    def execute(self, device: torch.device):
-        """Executes the job
-
-        :params:
-            device: torch device either CPUs or a specified GPU
-        """
-        self._execution_method(device)
-
-    def _execute_train(self, device: torch.device):
-        """
-        Executes training pipeline from scratch for a Job model.
-
-        :params:
-            device: torch device either CPUs or a specified GPU
-        """
-        self.optimizer.register_model_params(model_params=dict(self.model.named_parameters()))
-        self.lr_scheduler.register_optimizer(optimizer=self.optimizer)
-
-        self.trainer.set_num_epochs(num_epochs=self.num_epochs)
-
-        # initial evaluation
-        evaluation_results = self._evaluation_step(device)
-
-        # we store the initial model / last warmup model again
-        checkpointing_instruction = self.checkpointing_strategy.get_model_checkpoint_instruction(
-            num_epochs=self.num_epochs, current_epoch=self.current_epoch, evaluation_result=evaluation_results
-        )
-        self.run_checkpointing(checkpointing_instruction)
-
-        # if early stopping criterion is fulfilled we can stop the training progress
-        if self.early_stopping_strategy.is_stopping_criterion_fulfilled(
-            current_epoch=self.current_epoch, evaluation_results=evaluation_results
-        ):
-            return
-
-        # start the actual training loop
-        self.current_epoch += 1
-        self.trainer.set_current_epoch(self.current_epoch)
-        while not self.trainer.is_done():
-            self.logger.log(LogLevel.INFO, f"epoch: {self.current_epoch}")
-            self._train_step(device)
-            evaluation_results = self._evaluation_step(device)
-            self.lr_scheduler.step()
-
-            checkpointing_instruction = self.checkpointing_strategy.get_model_checkpoint_instruction(
-                num_epochs=self.num_epochs, current_epoch=self.current_epoch, evaluation_result=evaluation_results
-            )
-            self.run_checkpointing(checkpointing_instruction)
-            if self.early_stopping_strategy.is_stopping_criterion_fulfilled(
-                current_epoch=self.current_epoch, evaluation_results=evaluation_results
-            ):
-                break
-
-            self.current_epoch += 1
-
-    def _execute_warm_start(self, device: torch.device):
-        """
-        Executes training pipeline from last known checkpoint for a Job model.
-
-        :params:
-            device: torch device either CPUs or a specified GPU
-        """
-        if self.current_epoch > 0:
-            model_state = pickle.loads(
-                self.gs_api_client.get_checkpoint_resource(
-                    grid_search_id=self.grid_search_id,
-                    experiment_id=self.experiment_id,
-                    checkpoint_id=self.current_epoch,
-                    checkpoint_resource=CheckpointResource.model,
-                )
-            )
-            self.model.load_state_dict(model_state)
-
-            optimizer_state = pickle.loads(
-                self.gs_api_client.get_checkpoint_resource(
-                    grid_search_id=self.grid_search_id,
-                    experiment_id=self.experiment_id,
-                    checkpoint_id=self.current_epoch,
-                    checkpoint_resource=CheckpointResource.optimizer,
-                )
-            )
-            self.optimizer.load_state_dict(optimizer_state)
-
-            lr_scheduler_state = pickle.loads(
-                self.gs_api_client.get_checkpoint_resource(
-                    grid_search_id=self.grid_search_id,
-                    experiment_id=self.experiment_id,
-                    checkpoint_id=self.current_epoch,
-                    checkpoint_resource=CheckpointResource.lr_scheduler,
-                )
-            )
-            self.lr_scheduler.load_state_dict(lr_scheduler_state)
-
-            state_component_state = pickle.loads(
-                self.gs_api_client.get_checkpoint_resource(
-                    grid_search_id=self.grid_search_id,
-                    experiment_id=self.experiment_id,
-                    checkpoint_id=self.current_epoch,
-                    checkpoint_resource=CheckpointResource.stateful_components,
-                )
-            )
-            self.set_state(state_component_state)
-
-        self._execute_train(device)
-
-    def _execute_eval(self, device: torch.device):
-        """
-        Executes evalation pipeline for a Job model.
-
-        :params:
-            device: torch device either CPUs or a specified GPU
-        """
-        for epoch in self.num_epochs:
-            # TODO need to load model + stateful components for the respective epoch here
-            self._evaluation_step(device, epoch=epoch)
-
-
-class GymJobFactory:
-    @staticmethod
-    def get_gym_job(
-        grid_search_id: str,
-        experiment_id: int,
-        run_mode: RunMode,
-        num_epochs: int,
-        gs_api_client: GridSearchAPIClientIF,
-        experiment_status_logger: ExperimentStatusLogger = None,
-        warm_start_epoch: int = 0,
-        **components: Dict[str, Any],
-    ) -> AbstractGymJob:
-        return GymJob(
-            grid_search_id=grid_search_id,
-            experiment_id=experiment_id,
-            run_mode=run_mode,
-            num_epochs=num_epochs,
-            gs_api_client=gs_api_client,
-            experiment_status_logger=experiment_status_logger,
-            warm_start_epoch=warm_start_epoch,
-            **components,
-        )
diff --git a/src/ml_gym/gym/trainer.py b/src/ml_gym/gym/trainer.py
deleted file mode 100644
index e15b4e03..00000000
--- a/src/ml_gym/gym/trainer.py
+++ /dev/null
@@ -1,131 +0,0 @@
-from abc import abstractmethod
-from typing import Dict, List, Callable, Any
-from ml_gym.loss_functions.loss_functions import Loss
-from ml_gym.models.nn.net import NNModel
-from ml_gym.data_handling.dataset_loader import DatasetLoader
-import torch
-from ml_gym.batching.batch import InferenceResultBatch, DatasetBatch
-from ml_gym.gym.inference_component import InferenceComponent
-from ml_gym.gym.stateful_components import StatefulComponent
-from ml_gym.optimizers.optimizer import OptimizerAdapter
-import tqdm
-from ml_gym.util.logger import ConsoleLogger
-from ml_gym.gym.post_processing import PredictPostProcessingIF
-from ml_gym.error_handling.exception import ModelAlreadyFullyTrainedError
-
-
-class TrainComponent(StatefulComponent):
-    def __init__(self, inference_component: InferenceComponent, post_processors: List[PredictPostProcessingIF],
-                 loss_fun: Loss, show_progress: bool = False):
-        self.loss_fun = loss_fun
-        self.inference_component = inference_component
-        self.post_processors = post_processors
-        self.show_progress = show_progress
-        self.logger = ConsoleLogger("logger_train_component")
-
-    def train_batch(self, batch: DatasetBatch, model: NNModel, optimizer: OptimizerAdapter, device: torch.device):
-        model.zero_grad()
-        batch.to_device(device)
-        loss = self.calc_loss(model, batch)
-        loss.sum().backward()
-        optimizer.step()
-
-    def train_epoch(self, model: NNModel, optimizer: OptimizerAdapter, data_loader: DatasetLoader,
-                    device: torch.device, epoch: int, batch_processed_callback_fun: Callable = None) -> NNModel:
-        data_loader.device = device
-        self.map_batches(fun=self.train_batch,
-                         loader=data_loader,
-                         fun_params={"device": device,
-                                     "model": model,
-                                     "optimizer": optimizer},
-                         progress_info=f"Training {data_loader.dataset_name}  @epoch {epoch}",
-                         callback_fun=batch_processed_callback_fun)
-        return model
-
-    def forward_batch(self, dataset_batch: DatasetBatch, model: NNModel, device: torch.device,) -> InferenceResultBatch:
-        model = model.to(device)
-        dataset_batch.to_device(device)
-        inference_result_batch = self.inference_component.predict(model, dataset_batch, self.post_processors)
-        return inference_result_batch
-
-    def calc_loss(self, model: NNModel, batch: DatasetBatch) -> torch.Tensor:
-        forward_batch = self.inference_component.predict(model, batch)
-        loss = self.loss_fun(forward_batch)
-        return loss
-
-    @staticmethod
-    def map_batches(fun: Callable[[DatasetBatch, NNModel], Any], loader: DatasetLoader,
-                    fun_params: Dict[str, Any] = None, progress_info: str = None,
-                    callback_fun: Callable = None) -> List[InferenceResultBatch]:
-        """
-        Applies a function to each dataset_batch within a DatasetLoader
-        """
-        num_batches = len(loader)
-        processed_batches = 0
-        update_lag = max(1, int(num_batches/10))
-        fun_params = fun_params if fun_params is not None else dict()
-        result = []
-        if progress_info is not None:
-            for dataset_batch in tqdm.tqdm(loader, desc=progress_info):
-                result.append(fun(dataset_batch, **fun_params))
-                processed_batches += 1
-                if callback_fun is not None and (processed_batches % update_lag == 0 or processed_batches == num_batches):
-                    callback_fun(status="train",
-                                 num_batches=num_batches,
-                                 current_batch=processed_batches,
-                                 splits=[loader.dataset_tag],
-                                 current_split=loader.dataset_tag)
-        else:
-            for dataset_batch in loader:
-                result.append(fun(dataset_batch, **fun_params))
-                processed_batches += 1
-                if callback_fun is not None and (processed_batches % update_lag == 0 or processed_batches == num_batches):
-                    callback_fun(status="train",
-                                 num_batches=num_batches,
-                                 current_batch=processed_batches,
-                                 splits=[loader.dataset_tag],
-                                 current_split=loader.dataset_tag)
-        return result
-
-
-class TrainerIF(StatefulComponent):
-
-    @abstractmethod
-    def is_done(self) -> bool:
-        raise NotImplementedError
-
-    @abstractmethod
-    def set_num_epochs(self, num_epochs: int):
-        raise NotImplementedError
-
-    @abstractmethod
-    def set_current_epoch(self, epoch: int):
-        raise NotImplementedError
-
-
-class Trainer(TrainerIF):
-    def __init__(self, train_component: TrainComponent, train_loader: DatasetLoader, verbose=False):
-        self.train_component = train_component
-        self.train_loader = train_loader
-        self.verbose = verbose
-        self.current_epoch = 1
-        self.num_epochs = -1
-
-    def is_done(self) -> bool:
-        return self.current_epoch > self.num_epochs  # training starts at epoch 1, epoch 0 is just for evaluation, therefore >
-
-    def set_num_epochs(self, num_epochs: int):
-        self.num_epochs = num_epochs
-
-    def set_current_epoch(self, epoch: int):
-        self.current_epoch = epoch
-
-    def train_epoch(self, model: NNModel, optimizer: OptimizerAdapter, device: torch.device,
-                    batch_processed_callback_fun: Callable = None) -> NNModel:
-        if self.current_epoch > self.num_epochs:
-            raise ModelAlreadyFullyTrainedError(f"Model has been already trained for {self.current_epoch}/{self.num_epochs} epochs.")
-        self.train_loader.device = device
-        model = self.train_component.train_epoch(model, optimizer, self.train_loader, device, self.current_epoch,
-                                                 batch_processed_callback_fun)
-        self.current_epoch += 1
-        return model
diff --git a/src/ml_gym/gym/trainers/__init__.py b/src/ml_gym/gym/trainers/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/ml_gym/gym/trainers/accelerate_trainer.py b/src/ml_gym/gym/trainers/accelerate_trainer.py
new file mode 100644
index 00000000..913d58a8
--- /dev/null
+++ b/src/ml_gym/gym/trainers/accelerate_trainer.py
@@ -0,0 +1,97 @@
+from itertools import chain
+from typing import List, Callable
+from ml_gym.loss_functions.loss_functions import Loss
+from ml_gym.models.nn.net import NNModel
+from ml_gym.data_handling.dataset_loader import DatasetLoader
+import torch
+from ml_gym.batching.batch import DatasetBatch
+from ml_gym.gym.inference_component import InferenceComponent
+from ml_gym.gym.stateful_components import StatefulComponent
+from ml_gym.optimizers.optimizer import OptimizerAdapter
+from ml_gym.gym.post_processing import PredictPostProcessingIF
+from accelerate import Accelerator
+import numpy as np
+
+
+class AccelerateTrainComponent(StatefulComponent):
+    def __init__(self, inference_component: InferenceComponent, post_processors: List[PredictPostProcessingIF],
+                 loss_fun: Loss):
+        self.loss_fun = loss_fun
+        self.inference_component = inference_component
+        self.post_processors = post_processors
+
+    def _train_batch(self, accelerator: Accelerator, batch: DatasetBatch, model: NNModel, optimizer: OptimizerAdapter) -> NNModel:
+        model.zero_grad()
+        loss = self.calc_loss(model, batch).sum()
+
+        # if accelerator.is_main_process:
+        #     w = model.module.fc_layers[0].weight
+        #     print(f"\n\nBefore Update main thread: {w}")
+        #     accelerator.backward(loss)
+        #     optimizer.step()
+        #     w = model.module.fc_layers[0].weight
+        #     print(f"\n\nAfter Update main thread: {w}")
+        # else:
+        #     w = model.module.fc_layers[0].weight
+        #     print(f"\n\nBefore 2nd thread: {w}")
+        #     accelerator.backward(loss)
+        #     optimizer.step()
+        #     w = model.module.fc_layers[0].weight
+        #     print(f"\n\nAfter 2nd thread: {w}")
+        #     print("\n")
+
+        accelerator.backward(loss)
+        optimizer.step()
+
+        return model
+
+    def train(self, model: NNModel, optimizer: OptimizerAdapter, dataloader: DatasetLoader,
+              accelerator: Accelerator, batch_done_callback_fun: Callable, epoch_done_callback_fun: Callable,
+              num_epochs: int, num_batches_per_epoch: int = None) -> NNModel:
+
+        if num_batches_per_epoch is None:
+            num_batches_per_epoch = len(dataloader)
+
+        num_total_batches = num_batches_per_epoch*num_epochs
+        num_dataloaders = int(np.ceil(num_total_batches/len(dataloader)))
+        data_loaders = chain(*([dataloader]*num_dataloaders))
+
+        for batch_id, batch in zip(range(num_total_batches), data_loaders):
+            current_epoch = int(batch_id / num_batches_per_epoch)
+            model = self._train_batch(accelerator=accelerator, batch=batch, model=model, optimizer=optimizer)
+            if accelerator.is_main_process:
+                batch_done_callback_fun(status="train",
+                                        num_batches=num_batches_per_epoch,
+                                        current_batch=batch_id % num_batches_per_epoch,
+                                        splits=[dataloader.dataset_tag],
+                                        current_split=dataloader.dataset_tag,
+                                        num_epochs=num_epochs,
+                                        current_epoch=current_epoch)
+            if (batch_id + 1) % num_batches_per_epoch == 0:  # when epoch done
+                epoch_done_callback_fun(num_epochs=num_epochs, current_epoch=current_epoch, model=model, accelerator=accelerator)
+
+        return model
+
+    def calc_loss(self, model: NNModel, batch: DatasetBatch) -> torch.Tensor:
+        forward_batch = self.inference_component.predict(batch=batch, model=model, post_processors=self.post_processors)
+        loss = self.loss_fun(forward_batch)
+        return loss
+
+
+class AccelerateTrainer:
+    def __init__(self, train_component: AccelerateTrainComponent, train_loader: DatasetLoader):
+        self.train_component = train_component
+        self.train_loader = train_loader
+
+    def train(self, num_epochs: int, model: NNModel, optimizer: OptimizerAdapter,
+              batch_done_callback_fun: Callable, epoch_done_callback: Callable, accelerator: Accelerator,
+              num_batches_per_epoch: int = None) -> NNModel:
+
+        model = model.train()
+
+        # accelerate_train_loader = accelerator.prepare(self.train_loader)
+
+        model = self.train_component.train(model=model, optimizer=optimizer, dataloader=self.train_loader, accelerator=accelerator,
+                                           batch_done_callback_fun=batch_done_callback_fun, epoch_done_callback_fun=epoch_done_callback,
+                                           num_epochs=num_epochs, num_batches_per_epoch=num_batches_per_epoch)
+        return model
diff --git a/src/ml_gym/gym/trainers/standard_trainer.py b/src/ml_gym/gym/trainers/standard_trainer.py
new file mode 100644
index 00000000..e4361707
--- /dev/null
+++ b/src/ml_gym/gym/trainers/standard_trainer.py
@@ -0,0 +1,80 @@
+from itertools import chain
+from typing import List, Callable
+from ml_gym.loss_functions.loss_functions import Loss
+from ml_gym.models.nn.net import NNModel
+from ml_gym.data_handling.dataset_loader import DatasetLoader
+import torch
+from ml_gym.batching.batch import DatasetBatch
+from ml_gym.gym.inference_component import InferenceComponent
+from ml_gym.gym.stateful_components import StatefulComponent
+from ml_gym.optimizers.optimizer import OptimizerAdapter
+from ml_gym.gym.post_processing import PredictPostProcessingIF
+import numpy as np
+
+
+class TrainComponent(StatefulComponent):
+    def __init__(self, inference_component: InferenceComponent, post_processors: List[PredictPostProcessingIF],
+                 loss_fun: Loss):
+        self.loss_fun = loss_fun
+        self.inference_component = inference_component
+        self.post_processors = post_processors
+
+    def _train_batch(self, batch: DatasetBatch, model: NNModel, optimizer: OptimizerAdapter, device: torch.device):
+        model = model.to(device)
+        batch.to(device)
+        model.zero_grad()
+        loss = self.calc_loss(model, batch)
+        loss.sum().backward()
+        optimizer.step()
+        return model
+
+    def train(self, model: NNModel, optimizer: OptimizerAdapter, dataloader: DatasetLoader, device: torch.device,
+              batch_done_callback_fun: Callable, epoch_done_callback_fun: Callable,
+              num_epochs: int, num_batches_per_epoch: int = None) -> NNModel:
+
+        if num_batches_per_epoch is None:
+            num_batches_per_epoch = len(dataloader)
+
+        num_total_batches = num_batches_per_epoch*num_epochs
+        num_dataloaders = int(np.ceil(num_total_batches/len(dataloader)))
+        data_loaders = chain(*([dataloader]*num_dataloaders))
+        # train_bar = tqdm.tqdm(total=num_batches_per_epoch)
+
+        for batch_id, batch in zip(range(num_total_batches), data_loaders):
+            current_epoch = int(batch_id / num_batches_per_epoch)
+            model = self._train_batch(batch=batch, model=model, optimizer=optimizer, device=device)
+
+            batch_done_callback_fun(status="train",
+                                    num_batches=num_batches_per_epoch,
+                                    current_batch=batch_id % num_batches_per_epoch,
+                                    splits=[dataloader.dataset_tag],
+                                    current_split=dataloader.dataset_tag,
+                                    num_epochs=num_epochs,
+                                    current_epoch=current_epoch)
+
+            if (batch_id + 1) % num_batches_per_epoch == 0:  # when epoch done
+                epoch_done_callback_fun(num_epochs=num_epochs, current_epoch=current_epoch, model=model)
+
+        return model
+
+    def calc_loss(self, model: NNModel, batch: DatasetBatch) -> torch.Tensor:
+        forward_batch = self.inference_component.predict(batch=batch, model=model, post_processors=self.post_processors)
+        loss = self.loss_fun(forward_batch)
+        return loss
+
+
+class Trainer:
+    def __init__(self, train_component: TrainComponent, train_loader: DatasetLoader):
+        self.train_component = train_component
+        self.train_loader = train_loader
+
+    def train(self, num_epochs: int, model: NNModel, optimizer: OptimizerAdapter, device: torch.device,
+              batch_done_callback_fun: Callable, epoch_done_callback: Callable,
+              num_batches_per_epoch: int = None) -> NNModel:
+
+        model = model.train()
+
+        model = self.train_component.train(model=model, optimizer=optimizer, dataloader=self.train_loader, device=device,
+                                           batch_done_callback_fun=batch_done_callback_fun, epoch_done_callback_fun=epoch_done_callback,
+                                           num_epochs=num_epochs, num_batches_per_epoch=num_batches_per_epoch)
+        return model
diff --git a/src/ml_gym/gym/worker.py b/src/ml_gym/gym/worker.py
index 123b4e08..5f58a3e5 100644
--- a/src/ml_gym/gym/worker.py
+++ b/src/ml_gym/gym/worker.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from ml_gym.gym.jobs import AbstractGymJob
+from ml_gym.gym.gym_jobs.standard_gym_job import AbstractGymJob
 
 
 class WorkerIF(ABC):
diff --git a/src/ml_gym/io/websocket_client.py b/src/ml_gym/io/websocket_client.py
index a2c3ccae..6689a237 100644
--- a/src/ml_gym/io/websocket_client.py
+++ b/src/ml_gym/io/websocket_client.py
@@ -39,6 +39,9 @@ def connect(self):
         BufferedClient._register_callback_funs(self._sio_client)
         self.emit("join", {"client_id": self._client_id, "rooms": [*self.rooms, self._client_id]})
 
+    def disconnect(self):
+        self._sio_client.disconnect()
+
     def leave(self):
         self.emit("leave", None)
 
diff --git a/src/ml_gym/modes.py b/src/ml_gym/modes.py
index e7506434..f9f73484 100644
--- a/src/ml_gym/modes.py
+++ b/src/ml_gym/modes.py
@@ -9,5 +9,10 @@ class ValidationMode(Enum):
 
 class RunMode(Enum):
     TRAIN = "train"
-    # RE_EVAL = "re_eval"
     WARM_START = "warm_start"
+
+
+class ParallelizationMode(Enum):
+
+    PARALLEL_GS_SINGLE_NODE_MULTI_GPU = "PARALLEL_GS_SINGLE_NODE_MULTI_GPU"
+    SEQUENTIAL_GS_SINGLE_NODE_DISTRIBUTED = "SEQUENTIAL_GS_SINGLE_NODE_DISTRIBUTED"
diff --git a/src/ml_gym/multiprocessing/job.py b/src/ml_gym/multiprocessing/job.py
index 3e2208d5..898779ca 100644
--- a/src/ml_gym/multiprocessing/job.py
+++ b/src/ml_gym/multiprocessing/job.py
@@ -6,9 +6,6 @@
 
 
 class JobIF:
-    @property
-    def device(self) -> torch.device:
-        raise NotImplementedError
 
     @property
     def experiment_id(self) -> str:
@@ -18,32 +15,23 @@ def experiment_id(self) -> str:
     def grid_search_id(self) -> str:
         raise NotImplementedError
 
-    @device.setter
-    def device(self, value: torch.device):
-        raise NotImplementedError
-
     def execute(self):
         raise NotImplementedError
 
 
 class Job(JobIF):
-    def __init__(self, job_id: str, fun: Callable, blueprint: BluePrint, param_dict: Dict, job_type: JobType = JobType.CALC):
+    def __init__(self, job_id: str, fun: Callable, blueprint: BluePrint, param_dict: Dict = None, job_type: JobType = JobType.CALC):
         self.job_id = job_id
         self.job_type = job_type
         self.fun = fun
         self.blueprint = blueprint
-        self.param_dict = param_dict
+        self.param_dict = param_dict if param_dict is not None else {}
         self.status = JobStatus.INIT
         self.starting_time = -1
         self.finishing_time = -1
         self.executing_process_id = -1
         self.error = None
         self.stacktrace = None
-        self._device: torch.device = torch.device("cpu")
-
-    @property
-    def device(self) -> torch.device:
-        return self._device
 
     @property
     def experiment_id(self) -> str:
@@ -59,13 +47,11 @@ def grid_search_id(self) -> str:
         else:
             return None
 
-    @device.setter
-    def device(self, value: torch.device):
-        self._device = value
-
-    def execute(self):
-        self.param_dict["device"] = self._device
-        return self.fun(blueprint=self.blueprint, **self.param_dict)
+    def execute(self, device: torch.device = None):
+        if device is not None:
+            return self.fun(blueprint=self.blueprint, device=device, **self.param_dict)
+        else:
+            return self.fun(blueprint=self.blueprint, **self.param_dict)
 
 
 class JobStatusSubscriberIF:
diff --git a/src/ml_gym/multiprocessing/pool.py b/src/ml_gym/multiprocessing/pool.py
index 263a8662..35e716de 100644
--- a/src/ml_gym/multiprocessing/pool.py
+++ b/src/ml_gym/multiprocessing/pool.py
@@ -4,8 +4,6 @@
 from typing import List
 from ml_gym.multiprocessing.job import JobType, Job, JobCollection, JobStatusSubscriberIF
 from ml_gym.multiprocessing.worker import WorkerProcessWrapper
-from ml_gym.util.logger import QueuedLogging
-from ml_gym.util.logger import LogLevel, QLogger
 from ml_gym.multiprocessing.job import JobStatus
 from ml_gym.persistency.logging import JobStatusLogger, MLgymStatusLoggerCollectionConstructable
 
@@ -16,36 +14,33 @@ def __init__(self, logger: JobStatusLogger):
         self._logger = logger
 
     def callback_job_event(self, job: Job):
-        representation = {"job_id": job.job_id, "job_type": job.job_type, "grid_search_id": job.grid_search_id, "experiment_id": job.experiment_id, "status": job.status,
+        representation = {"job_id": job.job_id, "job_type": job.job_type, "grid_search_id": job.grid_search_id,
+                          "experiment_id": job.experiment_id, "status": job.status,
                           "starting_time": job.starting_time, "finishing_time": job.finishing_time, "error": job.error,
-                          "stacktrace": job.stacktrace, "device": job.device}
+                          "stacktrace": job.stacktrace}
         self._logger.log_job_status(**representation)
 
 
 class Pool:
-    def __init__(self, num_processes: int, devices: List[torch.device], max_jobs_per_process: int = 1,
-                 logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None):
+    def __init__(self, num_processes: int, devices: List[torch.device],
+                 logger_collection_constructable: MLgymStatusLoggerCollectionConstructable, max_jobs_per_process: int = 1):
         self.num_processes = num_processes
         self.job_q = Queue()
         self.job_update_q = Queue()
         self.devices = devices
         self.worker_processes = []
         self.max_jobs_per_process = max_jobs_per_process
-        self.logger: QLogger = QueuedLogging.get_qlogger("logger_pool")
-        self.logger.log(LogLevel.INFO, f"Initialized to run jobs on: {self.devices}")
+        logger_collection = logger_collection_constructable.construct()
+        self.job_status_logger = JobStatusLogger(logger=logger_collection)
+        subscriber = JobStatusLoggingSubscriber(self.job_status_logger)
         self.job_collection = JobCollection()
-        if logger_collection_constructable is not None:
-            logger_collection = logger_collection_constructable.construct()
-            job_status_logger = JobStatusLogger(logger=logger_collection)
-            subscriber = JobStatusLoggingSubscriber(job_status_logger)
-            self.job_collection.add_subscriber(subscriber)
+        self.job_collection.add_subscriber(subscriber)
 
     def add_job(self, job: Job):
         self.job_q.put(job)
         self.job_collection.add_or_update_job(job)
 
     def add_jobs(self, jobs: List[Job]):
-        self.logger.log(LogLevel.INFO, "Filling up job queue ... ")
         for job in tqdm.tqdm(jobs):
             self.job_q.put(job)
             self.job_collection.add_or_update_job(job)
@@ -54,30 +49,25 @@ def run(self):
         # we have to add the termination jobs at the end of the queue such that the processes stop working and don't get stuck in jobs_q.get()
         termination_jobs = [Job(job_id=i+len(self.job_collection), fun=None, blueprint=None, param_dict=None,
                                 job_type=JobType.TERMINATE) for i in range(self.num_processes)]
-        print(f"num_processes: {self.num_processes}")
         self.add_jobs(termination_jobs)
         # create and start worker processes
-        self.logger.log(LogLevel.INFO, f"Creating {self.num_processes} worker processes...")
-        for process_id in tqdm.tqdm(range(self.num_processes)):
+        for process_id in tqdm.tqdm(range(self.num_processes), desc="Creating processes"):
             self.create_or_replace_process(process_id, self.max_jobs_per_process)
-        self.logger.log(LogLevel.INFO, f"Starting {self.num_processes} worker processes...")
-        for p in tqdm.tqdm(self.worker_processes):
+        for p in tqdm.tqdm(self.worker_processes, desc="Starting processes"):
             p.start()
         # wait until all jobs are done
         while not self.job_collection.done:
             updated_job: Job = self.job_update_q.get()
             self.job_collection.add_or_update_job(updated_job)
             if updated_job.status == JobStatus.DONE and updated_job.error is not None:
-                self.logger.log(
-                    LogLevel.FATAL, f"FATAL! Job {updated_job.job_id} crashed while being executed by process {updated_job.executing_process_id} on {updated_job.device} after {int(updated_job.finishing_time - updated_job.starting_time)} seconds with error {updated_job.error}")
-                self.logger.log(LogLevel.FATAL, f"Error: {updated_job.stacktrace}")
+                pass
             elif updated_job.status == JobStatus.DONE:
-                self.logger.log(
-                    LogLevel.INFO, f"Job {updated_job.job_id} was successfully executed by process {updated_job.executing_process_id} on {updated_job.device} within {int(updated_job.finishing_time - updated_job.starting_time)} seconds.")
-                self.logger.log(LogLevel.INFO, f"Progress: {int(self.job_collection.done_count / self.job_collection.job_count * 100)}%")
+                pass
             if updated_job.status == JobStatus.DONE and updated_job.job_type == JobType.CALC:
                 self.worker_processes[updated_job.executing_process_id].recreate_process_if_done()
 
+        self.job_status_logger.disconnect()
+
     def create_or_replace_process(self, process_id: int, num_jobs_to_perform: int):
         process = WorkerProcessWrapper(process_id=process_id,
                                        num_jobs_to_perform=num_jobs_to_perform,
@@ -85,11 +75,8 @@ def create_or_replace_process(self, process_id: int, num_jobs_to_perform: int):
                                        job_q=self.job_q,
                                        job_update_q=self.job_update_q)
         if len(self.worker_processes) == process.process_id:
-            self.logger.log(LogLevel.DEBUG, f"Adding process {process_id}.")
             self.worker_processes.append(process)
         elif len(self.worker_processes) > process.process_id:
             self.worker_processes[process_id] = process
-            self.logger.log(LogLevel.DEBUG, f"Replaced process {process_id}.")
         else:
             raise Exception("There is something wrong with the process id!")
-
diff --git a/src/ml_gym/multiprocessing/worker.py b/src/ml_gym/multiprocessing/worker.py
index ea298a7e..a3318ac3 100644
--- a/src/ml_gym/multiprocessing/worker.py
+++ b/src/ml_gym/multiprocessing/worker.py
@@ -3,64 +3,60 @@
 import torch
 import traceback
 from ml_gym.multiprocessing.job import Job, JobType, JobStatus
-from ml_gym.util.logger import MLgymLoggerIF, LogLevel, QueuedLogging
 from copy import deepcopy
 
 
 class WorkerProcess(Process):
-    def __init__(self, process_id: int, num_jobs_to_perform: int, job_q: Queue, job_update_q: Queue, device: torch.device, logger: MLgymLoggerIF):
-        super(WorkerProcess, self).__init__(target=self.work, args=(job_q, job_update_q, num_jobs_to_perform, device, logger))
+    def __init__(self, process_id: int, num_jobs_to_perform: int, job_q: Queue, job_update_q: Queue, device: torch.device):
+        super(WorkerProcess, self).__init__(target=self.work, args=(job_q, job_update_q, num_jobs_to_perform, device))
         self.process_id = process_id
 
-    def work(self, job_q: Queue, job_update_q: Queue, num_jobs_to_perform: int, device: torch.device, logger: MLgymLoggerIF):
+    def work(self, job_q: Queue, job_update_q: Queue, num_jobs_to_perform: int, device: torch.device):
 
-        logger.log(LogLevel.INFO, f"Process {self.process_id} started working.")
         jobs_done_count = 0
-        for job in iter(job_q.get, None):  # https://stackoverflow.com/a/21157892
+        while True:
+            job: Job = job_q.get(block=True)            # https://stackoverflow.com/a/21157892
             job.status = JobStatus.RUNNING
-            job.device = device
             job.executing_process_id = self.process_id
-            logger.log(LogLevel.INFO, f"Process {job.executing_process_id} started job {job.job_id} on {job.device}.")
             job.starting_time = time.time()
             job_update_q.put(deepcopy(job))
             if job.job_type == JobType.CALC:
-                self._do_calc(job)
+                self._do_calc(job, device=device)
             job.finishing_time = time.time()
             job.status = JobStatus.DONE
             jobs_done_count += 1
             job_update_q.put(deepcopy(job))
             if job.job_type == JobType.TERMINATE or num_jobs_to_perform == jobs_done_count:
-                logger.log(LogLevel.DEBUG, f"Process {self.process_id} terminated.")
                 break
 
-    def _do_calc(self, job: Job):
+    def _do_calc(self, job: Job, device: torch.device):
         try:
-            job.execute()
+            job.execute(device=device)
         except Exception as e:
             job.error = str(e)
             job.stacktrace = traceback.format_exc()
+            print(f"ERROR: Job {job.job_id} running experiment {job.experiment_id} failed.")
+            print(job.error)
+            print(job.stacktrace)
 
 
 class WorkerProcessWrapper:
     def __init__(self, process_id: int, num_jobs_to_perform: int, device: torch.device, job_q: Queue, job_update_q: Queue):
-        self.logger = QueuedLogging.get_qlogger(f"logger_process_{process_id}")
         self.jobs_done_count = 0
         self.device = device
         self.num_jobs_to_perform = num_jobs_to_perform
         self.process_id = process_id
         self.job_q = job_q
         self.job_update_q = job_update_q
-        self.process = WorkerProcess(process_id, num_jobs_to_perform, job_q, job_update_q, device, self.logger)
+        self.process = WorkerProcess(process_id, num_jobs_to_perform, job_q, job_update_q, device)
 
     def recreate_process_if_done(self):
         self.jobs_done_count += 1
         if self.num_jobs_to_perform == self.jobs_done_count:
-            self.logger.log(LogLevel.DEBUG, f"Recreating process {self.process_id}.")
             self.process = WorkerProcess(self.process_id, self.num_jobs_to_perform,
-                                         self.job_q, self.job_update_q, self.device, self.logger)
+                                         self.job_q, self.job_update_q, self.device)
             self.jobs_done_count = 0
             self.process.start()
-            self.logger.log(LogLevel.DEBUG, f"Recreated process {self.process_id}.")
 
     def get_process_id(self) -> int:
         return self.process.process_id
diff --git a/src/ml_gym/optimizers/optimizer.py b/src/ml_gym/optimizers/optimizer.py
index 34757813..1646923e 100644
--- a/src/ml_gym/optimizers/optimizer.py
+++ b/src/ml_gym/optimizers/optimizer.py
@@ -4,7 +4,7 @@
 from ml_gym.error_handling.exception import OptimizerNotInitializedError
 
 
-class OptimizerAdapter(object):
+class OptimizerAdapter(Optimizer, object):
 
     def __init__(self, optimizer_class: Type[Optimizer], optimizer_params: Dict = None):
         self._optimizer_class = optimizer_class
diff --git a/src/ml_gym/persistency/io.py b/src/ml_gym/persistency/io.py
index a035cf0d..9493a1fd 100644
--- a/src/ml_gym/persistency/io.py
+++ b/src/ml_gym/persistency/io.py
@@ -258,17 +258,8 @@ def delete_checkpoint_resource(self, grid_search_id: str, experiment_id: str, ep
              experiment_id (str): Experiment ID
              epoch (int): Epoch
         """
-
-        chekpoint_resources = [
-            CheckpointResource.model,
-            CheckpointResource.lr_scheduler,
-            CheckpointResource.optimizer,
-            CheckpointResource.stateful_components,
-        ]
-
-        for checkpoint in chekpoint_resources:
-            url = f"{self.endpoint}/checkpoints/{grid_search_id}/{experiment_id}/{epoch}/{checkpoint}"
-            GridSearchRestfulAPIClient._del_binary_resource(url)
+        url = f"{self.endpoint}/checkpoints/{grid_search_id}/{experiment_id}/{epoch}"
+        GridSearchRestfulAPIClient._del_binary_resource(url)
 
     def get_experiment_statuses(self, grid_search_id: str) -> List[ExperimentStatus]:
         """
diff --git a/src/ml_gym/persistency/logging.py b/src/ml_gym/persistency/logging.py
index 2137403c..c2a22178 100644
--- a/src/ml_gym/persistency/logging.py
+++ b/src/ml_gym/persistency/logging.py
@@ -33,10 +33,14 @@ class MLgymStatusLoggerIF(ABC):
     def log_raw_message(self, raw_log_message: Dict):
         raise NotImplementedError
 
+    @abstractmethod
+    def disconnect(self):
+        raise NotImplementedError
+
 
 class LoggerCollection(MLgymStatusLoggerIF):
 
-    def __init__(self, loggers: MLgymStatusLoggerIF):
+    def __init__(self, loggers: List[MLgymStatusLoggerIF]):
         self._loggers = loggers
 
     # def log_job_status(self, job_id: int, job_type: JobType, status: JobStatus, experiment_id: str, starting_time: int, finishing_time: int, device: torch.device,
@@ -56,6 +60,10 @@ def log_raw_message(self, raw_log_message: Dict):
         for logger in self._loggers:
             logger.log_raw_message(raw_log_message)
 
+    def disconnect(self):
+        for logger in self._loggers:
+            logger.disconnect()
+
 
 class DiscLogger(MLgymStatusLoggerIF):
     def __init__(self):
@@ -80,6 +88,9 @@ def __init__(self, host: str, port: int):
     def log_raw_message(self, raw_log_message: Dict):
         self._sio_client.emit(message_key="mlgym_event", message=raw_log_message)
 
+    def disconnect(self):
+        self._sio_client.disconnect()
+
 
 class JobStatusLoggerIF(ABC):
 
@@ -90,16 +101,19 @@ def log_job_status(self, job_id: str, job_type: JobType, status: JobStatus, grid
     def log_experiment_config(self, grid_search_id: str, experiment_id: str, job_id: str, config: Dict[str, Any]):
         raise NotImplementedError
 
+    def disconnect(self):
+        raise NotImplementedError
+
 
 class JobStatusLogger(JobStatusLoggerIF):
     def __init__(self, logger: MLgymStatusLoggerIF) -> None:
         self._logger = logger
 
-    def log_job_status(self, job_id: str, job_type: JobType, status: JobStatus, grid_search_id: str, experiment_id: str, starting_time: int, finishing_time: int,
-                       device: torch.device, error: str = "", stacktrace: str = ""):
+    def log_job_status(self, job_id: str, job_type: JobType, status: JobStatus, grid_search_id: str, experiment_id: str,
+                       starting_time: int, finishing_time: int, error: str = "", stacktrace: str = ""):
         message = {"event_type": "job_status", "creation_ts": get_timestamp()}
-        payload = {"job_id": job_id, "job_type": job_type.value, "status": status.value, "grid_search_id": grid_search_id, "experiment_id": experiment_id,
-                   "starting_time": starting_time, "finishing_time": finishing_time, "device": str(device), "error": error,
+        payload = {"job_id": job_id, "job_type": job_type.value, "status": status.value, "grid_search_id": grid_search_id,
+                   "experiment_id": experiment_id, "starting_time": starting_time, "finishing_time": finishing_time, "error": error,
                    "stacktrace": stacktrace}
         message["payload"] = payload
         self._logger.log_raw_message(raw_log_message=message)
@@ -110,6 +124,9 @@ def log_experiment_config(self, grid_search_id: str, experiment_id: str, job_id:
         message["payload"] = payload
         self._logger.log_raw_message(raw_log_message=message)
 
+    def disconnect(self):
+        self._logger.disconnect()
+
 
 class ExperimentStatusLogger:
     def __init__(self, logger: MLgymStatusLoggerIF, experiment_id: str, grid_search_id: str, binary_stream_chunk_size: int = 1000000):
@@ -187,6 +204,9 @@ def get_chunks(binary_stream, binary_stream_chunk_size: int):
                 message["payload"] = payload
                 self._logger.log_raw_message(raw_log_message=message)
 
+    def disconnect(self):
+        self._logger.disconnect()
+
 
 class MLgymStatusLoggerTypes(Enum):
     DISC_LOGGER = DiscLogger
diff --git a/src/ml_gym/starter.py b/src/ml_gym/starter.py
deleted file mode 100644
index c4b26ae3..00000000
--- a/src/ml_gym/starter.py
+++ /dev/null
@@ -1,185 +0,0 @@
-import json
-from typing import List, Type
-from ml_board.backend.restful_api.data_models import FileFormat
-from ml_gym.blueprints.blue_prints import BluePrint
-from ml_gym.gym.gym import Gym
-from ml_gym.modes import RunMode
-from ml_gym.persistency.io import GridSearchAPIClientConstructableIF, GridSearchAPIClientIF
-from ml_gym.persistency.logging import LoggerConstructableIF, MLgymStatusLoggerCollectionConstructable
-from ml_gym.util.logger import QueuedLogging
-from ml_gym.validation.validator import ValidatorIF
-from datetime import datetime
-from ml_gym.io.config_parser import YAMLConfigLoader
-from ml_gym.error_handling.exception import SingletonAlreadyInstantiatedError
-
-
-class MLGymStarter:
-    @staticmethod
-    def _create_gym(job_id_prefix: str, process_count: int, device_ids, log_std_to_file: bool,
-                    logger_collection_constructable: MLgymStatusLoggerCollectionConstructable) -> Gym:
-        """
-        Creates Gym job Object
-
-        :params: 
-                job_id_prefix (str): unique ID
-                process_count (int): number of processors on which job will run
-                device_ids : GPU or CPU
-                log_std_to_file (bool): True or false based on config
-                logger_collection_constructable (MLgymStatusLoggerCollectionConstructable) : Websocket Server
-
-        :return: gym: Gym job object with config ready to be executed.
-
-        """
-        gym = Gym(job_id_prefix=job_id_prefix, process_count=process_count, device_ids=device_ids, log_std_to_file=log_std_to_file,
-                  logger_collection_constructable=logger_collection_constructable)
-        return gym
-
-    @staticmethod
-    def _setup_logging_environment(log_dir_path: str):
-        try:
-            QueuedLogging.start_logging(log_dir_path)
-        except SingletonAlreadyInstantiatedError:
-            pass
-
-    @staticmethod
-    def _stop_logging_environment():
-        QueuedLogging.stop_listener()
-
-
-# class MLGymTrainWarmStarter(MLGymStarter):
-
-#     def __init__(self, grid_search_id: str, blue_print_class: Type[BluePrint], num_epochs: int, text_logging_path: str, process_count: int,
-#                  gpus: List[int], log_std_to_file: bool, gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-#                  logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None) -> None:
-#         self.blue_print_class = blue_print_class
-#         self.num_epochs = num_epochs
-#         self.text_logging_path = text_logging_path
-#         self.process_count = process_count
-#         self.log_std_to_file = log_std_to_file
-#         self.gpus = gpus
-#         self.grid_search_id = grid_search_id
-#         self.gs_api_client_constructable = gs_api_client_constructable
-#         self.gs_api_client = gs_api_client_constructable.construct()
-#         self.logger_collection_constructable = logger_collection_constructable
-
-#     def start(self):
-#         self._setup_logging_environment(self.text_logging_path)
-#         gym = MLGymStarter._create_gym(process_count=self.process_count, device_ids=self.gpus, log_std_to_file=self.log_std_to_file,
-#                                        logger_collection_constructable=self.logger_collection_constructable)
-
-#         blueprints = self.validator.create_blueprints(blue_print_type=self.blue_print_class,
-#                                                       gs_config=self.gs_config,
-#                                                       num_epochs=self.num_epochs,
-#                                                       logger_collection_constructable=self.logger_collection_constructable)
-
-#         gym.add_blue_prints(blueprints)
-#         gym.run(parallel=True)
-
-#         self._stop_logging_environment()
-
-
-class MLGymTrainStarter(MLGymStarter):
-
-    def __init__(self, text_logging_path: str, process_count: int,
-                 gpus: List[int], log_std_to_file: bool, blueprints: List[BluePrint],
-                 logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None) -> None:
-        self.text_logging_path = text_logging_path
-        self.process_count = process_count
-        self.log_std_to_file = log_std_to_file
-        self.gpus = gpus
-        self.blueprints = blueprints
-        self.logger_collection_constructable = logger_collection_constructable
-
-    def start(self):
-        self._setup_logging_environment(self.text_logging_path)
-        job_id_prefix = datetime.now().strftime("%Y-%m-%d--%H-%M-%S")
-
-        gym = MLGymStarter._create_gym(job_id_prefix=job_id_prefix, process_count=self.process_count, device_ids=self.gpus,
-                                       log_std_to_file=self.log_std_to_file,
-                                       logger_collection_constructable=self.logger_collection_constructable)
-        gym.add_blueprints(self.blueprints)
-        gym.run(parallel=True)
-
-        self._stop_logging_environment()
-
-
-def save_blueprint_config(blueprint: BluePrint, gs_api_client: GridSearchAPIClientIF):
-    gs_api_client.add_config_string(grid_search_id=blueprint.grid_search_id, config_name="experiment_config.json",
-                                    config=json.dumps(blueprint.config), experiment_id=blueprint.experiment_id,
-                                    file_format=FileFormat.JSON)
-
-
-def mlgym_entry_train(blueprint_class: Type[BluePrint],
-                      logger_collection_constructable: LoggerConstructableIF,
-                      gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-                      gs_config_raw_string: str,
-                      validator: ValidatorIF,
-                      text_logging_path: str,
-                      process_count: int,
-                      gpus: int,
-                      log_std_to_file: bool,
-                      num_epochs: int,
-                      validation_strategy_config_raw_string: str = None):
-    gs_api_client = gs_api_client_constructable.construct()
-
-    grid_search_id = datetime.now().strftime("%Y-%m-%d--%H-%M-%S")
-    gs_config = YAMLConfigLoader.load_string(gs_config_raw_string)
-    gs_api_client.add_config_string(grid_search_id=grid_search_id,
-                                    file_format=FileFormat.YAML,
-                                    config_name="grid_search_config.yml",
-                                    config=gs_config_raw_string)
-
-    if validation_strategy_config_raw_string is not None:
-        gs_api_client.add_config_string(grid_search_id=grid_search_id,
-                                        config_name="validation_strategy_config.yml",
-                                        config=validation_strategy_config_raw_string)
-
-    blueprints = validator.create_blueprints(grid_search_id=grid_search_id,
-                                             blue_print_type=blueprint_class,
-                                             gs_config=gs_config,
-                                             num_epochs=num_epochs,
-                                             gs_api_client_constructable=gs_api_client_constructable,
-                                             logger_collection_constructable=logger_collection_constructable)
-
-    for blueprint in blueprints:
-        save_blueprint_config(blueprint=blueprint, gs_api_client=gs_api_client)
-
-    starter = MLGymTrainStarter(text_logging_path=text_logging_path,
-                                process_count=process_count,
-                                gpus=gpus,
-                                log_std_to_file=log_std_to_file,
-                                blueprints=blueprints,
-                                logger_collection_constructable=logger_collection_constructable)
-    starter.start()
-
-
-def mlgym_entry_warm_start(blueprint_class: Type[BluePrint],
-                           grid_search_id: str,
-                           logger_collection_constructable: LoggerConstructableIF,
-                           gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-                           text_logging_path: str,
-                           process_count: int,
-                           gpus: int,
-                           log_std_to_file: bool,
-                           num_epochs: int):
-    gs_api_client = gs_api_client_constructable.construct()
-    experiment_statuses = gs_api_client.get_experiment_statuses(grid_search_id)
-
-    blueprints = [BluePrint.create_blueprint(blue_print_class=blueprint_class,
-                                             run_mode=RunMode.WARM_START,
-                                             experiment_config=experiment_status.experiment_config,
-                                             num_epochs=num_epochs,
-                                             warm_start_epoch=experiment_status.last_checkpoint_id,
-                                             grid_search_id=grid_search_id,
-                                             experiment_id=experiment_status.experiment_id,
-                                             logger_collection_constructable=logger_collection_constructable,
-                                             gs_api_client_constructable=gs_api_client_constructable)
-                  for experiment_status in experiment_statuses if experiment_status.last_checkpoint_id < num_epochs]
-
-    starter = MLGymTrainStarter(text_logging_path=text_logging_path,
-                                process_count=process_count,
-                                gpus=gpus,
-                                log_std_to_file=log_std_to_file,
-                                blueprints=blueprints,
-                                logger_collection_constructable=logger_collection_constructable)
-    starter.start()
diff --git a/src/ml_gym/util/logger.py b/src/ml_gym/util/logger.py
deleted file mode 100644
index 7c439db0..00000000
--- a/src/ml_gym/util/logger.py
+++ /dev/null
@@ -1,169 +0,0 @@
-import logging
-import os
-from multiprocessing import Queue, Process
-from dataclasses import dataclass
-import time
-from abc import ABC, abstractmethod
-from enum import Enum
-import sys
-import traceback
-from ml_gym.error_handling.exception import SingletonAlreadyInstantiatedError
-
-
-class LogLevel(Enum):
-    CRITICAL = 50
-    FATAL = CRITICAL
-    ERROR = 40
-    WARNING = 30
-    WARN = WARNING
-    INFO = 20
-    DEBUG = 10
-    NOTSET = 0
-
-
-@dataclass
-class Message:
-    logger_id: str
-    message_string: str
-    unix_timestamp: float
-    level: LogLevel
-
-
-class MLgymLoggerIF(ABC):
-
-    def __init__(self, logger_id: str):
-        self.logger_id = logger_id
-
-    def build_message(self, level: LogLevel, message) -> Message:
-        timestamp = time.time()
-        message = Message(message_string=message, unix_timestamp=timestamp, level=level, logger_id=self.logger_id)
-        return message
-
-    @abstractmethod
-    def log(self, level: LogLevel, message: str):
-        raise NotImplementedError
-
-
-class ConsoleLogger(MLgymLoggerIF):
-
-    def __init__(self, logger_id: str):
-        super().__init__(logger_id)
-        self.logger = ConsoleLogger._get_console_logger(logger_id)
-
-    def log(self, level: LogLevel, message: str):
-        self.logger.log(level.value, message)
-
-    @staticmethod
-    def _get_console_logger(logger_id: str):
-        # create / get logger
-        logger = logging.getLogger(logger_id)
-        # if logger has not been created
-        if not logger.hasHandlers():
-            logger.setLevel(logging.DEBUG)
-
-            # create console handler with a higher log level
-            ch = logging.StreamHandler()
-            ch.setLevel(logging.DEBUG)
-            # create formatter and add it to the handlers
-            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-            ch.setFormatter(formatter)
-            # add the handlers to the logger
-            logger.addHandler(ch)
-        return logger
-
-
-class QLogger(MLgymLoggerIF):
-
-    def __init__(self, logger_id: str, queue: Queue):
-        super().__init__(logger_id)
-        self.queue = queue
-
-    def log(self, level: LogLevel, message: str):
-        message = self.build_message(level, message)
-        self.queue.put(message)
-
-
-class QueuedLogging:
-    _instance: "QueuedLogging" = None
-
-    @classmethod
-    def start_logging(cls, log_dir_path: str):
-        if cls._instance is None:
-            log_msg_queue = Queue()
-            cls._instance = QueuedLogging(log_msg_queue, log_dir_path)
-            cls._instance._start_listening()
-        else:
-            raise SingletonAlreadyInstantiatedError
-
-    @classmethod
-    def get_qlogger(cls, logger_id: str) -> QLogger:
-        return QLogger(logger_id, QueuedLogging._instance.log_msg_queue)
-
-    def __init__(self, log_msg_queue: Queue, log_dir_path: str):
-        if self._instance is not None:
-            raise SingletonAlreadyInstantiatedError
-        self.log_msg_queue = log_msg_queue
-        self.log_dir_path = log_dir_path
-        self.listener_process = None
-
-    def _start_listening(self):
-        self.listener_process = Process(target=QueuedLogging._listener_process, args=(self.log_msg_queue, self.log_dir_path))
-        self.listener_process.start()
-
-    @staticmethod
-    def stop_listener():
-        QueuedLogging._instance.log_msg_queue.put_nowait(None)
-
-    @staticmethod
-    def _get_logger(log_dir_path: str, logger_id: str):
-        logger = logging.getLogger(logger_id)
-        log_file_path = os.path.join(log_dir_path, f"{logger_id}.log")
-        os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
-        logger.setLevel(logging.DEBUG)
-        # create file handler which logs even debug messages
-        fh = logging.FileHandler(log_file_path)
-        fh.setLevel(logging.DEBUG)
-        # create console handler with a higher log level
-        ch = logging.StreamHandler()
-        ch.setLevel(logging.DEBUG)
-        # create formatter and add it to the handlers
-        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-        fh.setFormatter(formatter)
-        ch.setFormatter(formatter)
-        # add the handlers to the logger
-        logger.addHandler(fh)
-        logger.addHandler(ch)
-        return logger
-
-    @staticmethod
-    def _listener_process(queue: Queue, log_dir_path: str):
-        loggers = {}
-        while True:
-            try:
-                message: Message = queue.get()
-                if message is None:  # We send this as a sentinel to tell the listener to quit.
-                    print("Logging listener process stopped.")
-                    break
-
-                if message.logger_id not in loggers:
-                    loggers[message.logger_id] = QueuedLogging._get_logger(log_dir_path, message.logger_id)
-                loggers[message.logger_id].log(level=message.level.value, msg=message.message_string)
-            except Exception:
-                print('Whoops! Problem:', file=sys.stderr)
-                traceback.print_exc(file=sys.stderr)
-
-
-def get_console_logger(name: str):
-    # create logger with 'spam_application'
-    logger = logging.getLogger(name)
-    logger.setLevel(logging.DEBUG)
-
-    # create console handler with a higher log level
-    ch = logging.StreamHandler()
-    ch.setLevel(logging.DEBUG)
-    # create formatter and add it to the handlers
-    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-    ch.setFormatter(formatter)
-    # add the handlers to the logger
-    logger.addHandler(ch)
-    return logger
diff --git a/src/ml_gym/util/util.py b/src/ml_gym/util/util.py
index fbea0d11..9ae165f7 100644
--- a/src/ml_gym/util/util.py
+++ b/src/ml_gym/util/util.py
@@ -15,7 +15,7 @@
 from ml_gym.gym.predict_postprocessing_component import PredictPostprocessingComponent
 from ml_gym.gym.post_processing import PredictPostProcessingIF
 import tqdm
-from ml_gym.gym.jobs import AbstractGymJob
+from ml_gym.gym.gym_jobs.standard_gym_job import AbstractGymJob
 from data_stack.dataset.iterator import InformedDatasetIteratorIF
 
 
@@ -48,7 +48,7 @@ def predict_tensor(self, sample_tensor: torch.Tensor, targets: torch.Tensor = No
         return result_batch
 
     def predict_dataset_batch(self, batch: DatasetBatch, no_grad: bool = True) -> InferenceResultBatch:
-        batch.to_device(self._device)
+        batch.to(self._device)
         if no_grad:
             with torch.no_grad():
                 forward_result = self.model.forward(batch.samples)
diff --git a/src/ml_gym/validation/cross_validation.py b/src/ml_gym/validation/cross_validation.py
index 49e19e29..e237d453 100644
--- a/src/ml_gym/validation/cross_validation.py
+++ b/src/ml_gym/validation/cross_validation.py
@@ -3,11 +3,9 @@
 from ml_gym.blueprints.blue_prints import BluePrint
 from data_stack.dataset.splitter import SplitterFactory
 from ml_gym.modes import RunMode
-from ml_gym.persistency.logging import MLgymStatusLoggerCollectionConstructable
 from ml_gym.validation.validator import ValidatorIF
 from ml_gym.blueprints.component_factory import Injector
 from ml_gym.util.grid_search import GridSearch
-from ml_gym.persistency.io import GridSearchAPIClientConstructableIF
 
 
 class CrossValidation(ValidatorIF):
@@ -52,9 +50,7 @@ def _create_folds_splits(folds_indices: List[List[int]]) -> List[Dict[str, Any]]
             splits.append(split)
         return splits
 
-    def create_blue_prints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any], num_epochs: int,
-                           gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-                           logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None) -> List[Type[BluePrint]]:
+    def create_blue_prints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any]) -> List[Type[BluePrint]]:
 
         run_id_to_config_dict = {run_id: config for run_id, config in enumerate(GridSearch.create_gs_from_config_dict(gs_config))}
 
@@ -75,20 +71,12 @@ def create_blue_prints(self, grid_search_id: str, blue_print_type: Type[BluePrin
                                                 run_mode=self.run_mode,
                                                 experiment_id=experiment_id,
                                                 experiment_config=experiment_config_injected,
-                                                num_epochs=num_epochs,
-                                                grid_search_id=grid_search_id,
-                                                logger_collection_constructable=logger_collection_constructable,
-                                                gs_api_client_constructable=gs_api_client_constructable)
+                                                grid_search_id=grid_search_id)
                 blueprints.append(bp)
                 experiment_id = experiment_id + 1
         return blueprints
 
-    def create_blueprints(self, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any], num_epochs: int,
-                          gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-                          logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None) -> List[BluePrint]:
+    def create_blueprints(self, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any]) -> List[BluePrint]:
         blueprints = self.create_blue_prints(blue_print_type=blue_print_type,
-                                             gs_config=gs_config,
-                                             num_epochs=num_epochs,
-                                             logger_collection_constructable=logger_collection_constructable,
-                                             gs_api_client_constructable=gs_api_client_constructable)
+                                             gs_config=gs_config)
         return blueprints
diff --git a/src/ml_gym/validation/gs_validator.py b/src/ml_gym/validation/gs_validator.py
index 2076ff9c..ca9bc8e0 100644
--- a/src/ml_gym/validation/gs_validator.py
+++ b/src/ml_gym/validation/gs_validator.py
@@ -1,19 +1,15 @@
 from typing import Dict, Any, Type, List
 from ml_gym.blueprints.blue_prints import BluePrint
 from ml_gym.modes import RunMode
-from ml_gym.persistency.logging import MLgymStatusLoggerCollectionConstructable
 from ml_gym.validation.validator import ValidatorIF
 from ml_gym.util.grid_search import GridSearch
-from ml_gym.persistency.io import GridSearchAPIClientConstructableIF
 
 
 class GridSearchValidator(ValidatorIF):
     def __init__(self, run_mode: RunMode):
         self.run_mode = run_mode
 
-    def create_blueprints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any], num_epochs: int,
-                          gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-                          logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None) -> List[BluePrint]:
+    def create_blueprints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any]) -> List[BluePrint]:
         run_id_to_config_dict = {run_id: config for run_id, config in enumerate(GridSearch.create_gs_from_config_dict(gs_config))}
 
         blueprints = []
@@ -21,10 +17,7 @@ def create_blueprints(self, grid_search_id: str, blue_print_type: Type[BluePrint
             blueprint = BluePrint.create_blueprint(blue_print_class=blue_print_type,
                                                    run_mode=self.run_mode,
                                                    experiment_config=experiment_config,
-                                                   num_epochs=num_epochs,
                                                    grid_search_id=grid_search_id,
-                                                   experiment_id=config_id,
-                                                   logger_collection_constructable=logger_collection_constructable,
-                                                   gs_api_client_constructable=gs_api_client_constructable)
+                                                   experiment_id=config_id)
             blueprints.append(blueprint)
         return blueprints
diff --git a/src/ml_gym/validation/nested_cross_validation.py b/src/ml_gym/validation/nested_cross_validation.py
index 907b948b..a794e3b8 100644
--- a/src/ml_gym/validation/nested_cross_validation.py
+++ b/src/ml_gym/validation/nested_cross_validation.py
@@ -3,11 +3,9 @@
 from ml_gym.blueprints.blue_prints import BluePrint
 from data_stack.dataset.splitter import SplitterFactory
 from ml_gym.modes import RunMode
-from ml_gym.persistency.logging import MLgymStatusLoggerCollectionConstructable
 from ml_gym.validation.validator import ValidatorIF
 from ml_gym.blueprints.component_factory import Injector
 from ml_gym.util.grid_search import GridSearch
-from ml_gym.persistency.io import GridSearchAPIClientConstructableIF
 
 
 class NestedCV(ValidatorIF):
@@ -81,8 +79,7 @@ def _create_inner_folds_splits(inner_folds_indices: Tuple[List[int]]) -> List[Di
                 splits.append(split)
         return splits
 
-    def _get_blue_prints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any], num_epochs: int,
-                         logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None) -> List[Type[BluePrint]]:
+    def _get_blue_prints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any]) -> List[Type[BluePrint]]:
 
         run_id_to_config_dict = {run_id: config for run_id, config in enumerate(GridSearch.create_gs_from_config_dict(gs_config))}
 
@@ -104,21 +101,14 @@ def _get_blue_prints(self, grid_search_id: str, blue_print_type: Type[BluePrint]
                 bp = BluePrint.create_blueprint(blue_print_class=blue_print_type,
                                                 run_mode=self.run_mode,
                                                 experiment_config=experiment_config_injected,
-                                                num_epochs=num_epochs,
                                                 grid_search_id=grid_search_id,
-                                                experiment_id=experiment_id,
-                                                logger_collection_constructable=logger_collection_constructable)
+                                                experiment_id=experiment_id)
                 blueprints.append(bp)
                 experiment_id = experiment_id + 1
         return blueprints
 
-    def create_blueprints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any], num_epochs: int,
-                          gs_api_client_constructable: GridSearchAPIClientConstructableIF,
-                          logger_collection_constructable: MLgymStatusLoggerCollectionConstructable = None) -> List[BluePrint]:
+    def create_blueprints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any]) -> List[BluePrint]:
         blueprints = self._get_blue_prints(grid_search_id=grid_search_id,
                                            blue_print_type=blue_print_type,
-                                           gs_config=gs_config,
-                                           num_epochs=num_epochs,
-                                           logger_collection_constructable=logger_collection_constructable,
-                                           gs_api_client_constructable=gs_api_client_constructable)
+                                           gs_config=gs_config)
         return blueprints
diff --git a/src/ml_gym/validation/validator.py b/src/ml_gym/validation/validator.py
index 35d9597c..1dfb0968 100644
--- a/src/ml_gym/validation/validator.py
+++ b/src/ml_gym/validation/validator.py
@@ -1,13 +1,10 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Any, Type, List
 from ml_gym.blueprints.blue_prints import BluePrint
-from ml_gym.persistency.io import GridSearchAPIClientConstructableIF
 
 
 class ValidatorIF(ABC):
 
     @abstractmethod
-    def create_blueprints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any],
-                          num_epochs: int, dashify_logging_path: str,
-                          gs_api_client_constructable: GridSearchAPIClientConstructableIF) -> List[BluePrint]:
+    def create_blueprints(self, grid_search_id: str, blue_print_type: Type[BluePrint], gs_config: Dict[str, Any]) -> List[BluePrint]:
         raise NotImplementedError
diff --git a/src/ml_gym/validation/validator_factory.py b/src/ml_gym/validation/validator_factory.py
index 81e7aaee..00d8131e 100644
--- a/src/ml_gym/validation/validator_factory.py
+++ b/src/ml_gym/validation/validator_factory.py
@@ -7,6 +7,7 @@
 from ml_gym.error_handling.exception import ValidationModeNotValidError
 import torch
 
+
 class ValidatorFactory:
 
     @staticmethod
diff --git a/src/setup.py b/src/setup.py
index 04247e23..018c17f8 100644
--- a/src/setup.py
+++ b/src/setup.py
@@ -32,6 +32,8 @@
         "eventlet",
         "python-socketio[client]",
         "gunicorn",
+        "flask-socketio",
+        "datasets",
         "transformers"
     ],
     python_requires=">=3.7",


Test patch:
diff --git a/pytests/batching/test_batch.py b/pytests/batching/test_batch.py
index 9bb82b1b..95980e82 100644
--- a/pytests/batching/test_batch.py
+++ b/pytests/batching/test_batch.py
@@ -22,13 +22,13 @@ def inference_batch_result(self) -> InferenceResultBatch:
 
     # test device operations
     def test_to_device(self, inference_batch_result: InferenceResultBatch):
-        inference_batch_result.to_device(torch.device("cpu"))
+        inference_batch_result.to(torch.device("cpu"))
 
     def test_to_cpu(self, inference_batch_result: InferenceResultBatch):
         inference_batch_result.to_cpu()
 
     def test_get_device(self, inference_batch_result: InferenceResultBatch):
-        inference_batch_result.to_device(torch.device("cpu"))
+        inference_batch_result.to(torch.device("cpu"))
         assert inference_batch_result.get_device() == torch.device("cpu")
 
     def test_detach(self, inference_batch_result: InferenceResultBatch):
@@ -61,13 +61,13 @@ def dataset_batch(self) -> DatasetBatch:
 
     # test device operations
     def test_to_device(self, dataset_batch: DatasetBatch):
-        dataset_batch.to_device(torch.device("cpu"))
+        dataset_batch.to(torch.device("cpu"))
 
     def test_to_cpu(self, dataset_batch: DatasetBatch):
         dataset_batch.to_cpu()
 
     def test_get_device(self, dataset_batch: DatasetBatch):
-        dataset_batch.to_device(torch.device("cpu"))
+        dataset_batch.to(torch.device("cpu"))
         assert dataset_batch.get_device() == torch.device("cpu")
 
     def test_detach(self, dataset_batch: DatasetBatch):
diff --git a/pytests/blueprints/constructables/test_eval_component_constructable.py b/pytests/blueprints/constructables/test_eval_component_constructable.py
index a976b076..e96e7b40 100644
--- a/pytests/blueprints/constructables/test_eval_component_constructable.py
+++ b/pytests/blueprints/constructables/test_eval_component_constructable.py
@@ -3,7 +3,7 @@
 from typing import List, Dict
 import pytest
 from ml_gym.blueprints.constructables import Requirement, EvalComponentConstructable
-from ml_gym.gym.evaluator import EvalComponent
+from ml_gym.gym.evaluators.evaluator import EvalComponent
 from pytests.blueprints.constructables.test_train_component_constructable import RegistryFixture, DataLoaderFixture
 
 
diff --git a/pytests/blueprints/constructables/test_evaluator_constructable.py b/pytests/blueprints/constructables/test_evaluator_constructable.py
index cbad476b..7c627342 100644
--- a/pytests/blueprints/constructables/test_evaluator_constructable.py
+++ b/pytests/blueprints/constructables/test_evaluator_constructable.py
@@ -2,7 +2,7 @@
 # -*- coding: UTF-8 -*-
 import pytest
 from ml_gym.blueprints.constructables import Requirement, EvalComponentConstructable, EvaluatorConstructable
-from ml_gym.gym.evaluator import Evaluator
+from ml_gym.gym.evaluators.evaluator import Evaluator
 from pytests.blueprints.constructables.test_eval_component_constructable import EvalConfigFixture
 from pytests.blueprints.constructables.test_train_component_constructable import RegistryFixture, DataLoaderFixture
 
diff --git a/pytests/blueprints/constructables/test_train_component_constructable.py b/pytests/blueprints/constructables/test_train_component_constructable.py
index 47cbe450..f7569fa1 100644
--- a/pytests/blueprints/constructables/test_train_component_constructable.py
+++ b/pytests/blueprints/constructables/test_train_component_constructable.py
@@ -6,7 +6,7 @@
     PredictionPostProcessingRegistryConstructable, Requirement,  DataLoadersConstructable, \
     MetricFunctionRegistryConstructable
 from ml_gym.gym.post_processing import ArgmaxPostProcessorImpl
-from ml_gym.gym.trainer import TrainComponent
+from ml_gym.gym.trainers.standard_trainer import TrainComponent
 from ml_gym.loss_functions.loss_functions import CrossEntropyLoss
 from pytests.blueprints.constructables.test_data_loader_constructable import CollatorFixture
 
diff --git a/pytests/blueprints/constructables/test_trainer_constructable.py b/pytests/blueprints/constructables/test_trainer_constructable.py
index 37984502..42c97cc7 100644
--- a/pytests/blueprints/constructables/test_trainer_constructable.py
+++ b/pytests/blueprints/constructables/test_trainer_constructable.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: UTF-8 -*-
 import pytest
-from ml_gym.gym.trainer import Trainer
+from ml_gym.gym.trainers.standard_trainer import Trainer
 from pytests.blueprints.constructables.test_train_component_constructable import RegistryFixture, DataLoaderFixture, TrainConfigFixture
 from ml_gym.blueprints.constructables import Requirement, TrainerConstructable,  \
     TrainComponentConstructable
diff --git a/pytests/gym/test_trainer.py b/pytests/gym/test_trainer.py
index a3e8e07f..26af9a8a 100644
--- a/pytests/gym/test_trainer.py
+++ b/pytests/gym/test_trainer.py
@@ -6,7 +6,7 @@
 from ml_gym.data_handling.dataset_loader import DatasetLoader
 from ml_gym.gym.inference_component import InferenceComponent
 from ml_gym.gym.post_processing import PredictPostProcessingIF
-from ml_gym.gym.trainer import TrainComponent, Trainer
+from ml_gym.gym.trainers.standard_trainer import TrainComponent, Trainer
 from ml_gym.loss_functions.loss_functions import Loss
 from ml_gym.models.nn.net import NNModel
 from ml_gym.optimizers.optimizer import OptimizerAdapter
diff --git a/pytests/test_env/conv_net_blueprint.py b/pytests/test_env/conv_net_blueprint.py
index 3ce2d489..0bbb8594 100644
--- a/pytests/test_env/conv_net_blueprint.py
+++ b/pytests/test_env/conv_net_blueprint.py
@@ -1,7 +1,7 @@
 from typing import Dict, List, Any
 from ml_gym.blueprints.constructables import ModelRegistryConstructable
 from ml_gym.blueprints.blue_prints import BluePrint
-from ml_gym.gym.jobs import AbstractGymJob, GymJobFactory
+from ml_gym.gym.gym_jobs.standard_gym_job import AbstractGymJob, GymJobFactory
 from ml_gym.batching.batch import DatasetBatch
 from dataclasses import dataclass
 from ml_gym.blueprints.component_factory import ComponentFactory, Injector
diff --git a/pytests/test_env/linear_net_blueprint.py b/pytests/test_env/linear_net_blueprint.py
index 11b82ba8..b4f32138 100644
--- a/pytests/test_env/linear_net_blueprint.py
+++ b/pytests/test_env/linear_net_blueprint.py
@@ -1,4 +1,4 @@
-import os
+wimport os
 from typing import List, Any, Tuple
 from data_stack.dataset.factory import BaseDatasetFactory
 from data_stack.dataset.iterator import InformedDatasetIteratorIF, SequenceDatasetIterator, \
@@ -8,7 +8,7 @@
 from ml_gym.blueprints.constructables import ModelRegistryConstructable, ComponentConstructable
 from ml_gym.blueprints.blue_prints import BluePrint
 from ml_gym.data_handling.postprocessors.factory import ModelGymInformedIteratorFactory
-from ml_gym.gym.jobs import AbstractGymJob, GymJobFactory
+from ml_gym.gym.gym_jobs.standard_gym_job import AbstractGymJob, GymJobFactory
 from ml_gym.batching.batch import DatasetBatch
 from dataclasses import dataclass, field
 from ml_gym.blueprints.component_factory import ComponentFactory, Injector