[feat] Add modular configs for slurm and aim. Decouple hf datasets. Better namings

Hovhannes Tamoyan · Hovhannes Tamoyan · commit 2cf9f5dafefc · 2024-07-18T11:44:59.000+02:00
diff --git a/urartu/VERSION b/urartu/VERSION
@@ -1 +1 @@
-2.0.22
+2.0.23
diff --git a/urartu/commands/__init__.py b/urartu/commands/__init__.py
@@ -1,7 +1,7 @@
 import argparse
 from typing import Tuple
 
-from .command import Command
+from urartu.commands.command import Command
 
 from .register import Register
 from .unregister import Unregister
diff --git a/urartu/commands/command.py b/urartu/commands/command.py
@@ -3,7 +3,7 @@
 
 from overrides import overrides
 
-from ..utils.registrable import Registrable
+from urartu.utils.registrable import Registrable
 
 T = TypeVar("T", bound="Command")
 
diff --git a/urartu/commands/launch.py b/urartu/commands/launch.py
@@ -6,11 +6,11 @@
 from hydra import compose, initialize
 from hydra.core.plugins import Plugins
 from omegaconf import OmegaConf
-from .command import Command
-from ..utils.launcher import launch, launch_on_slurm
-from ..utils.slurm import is_submitit_available
-from ..utils.registry import Registry
-from ..utils.hydra_plugin import UrartuPlugin
+from urartu.commands.command import Command
+from urartu.utils.launcher import launch, launch_on_slurm
+from urartu.utils.slurm import is_submitit_available
+from urartu.utils.registry import Registry
+from urartu.utils.hydra_plugin import UrartuPlugin
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
diff --git a/urartu/commands/register.py b/urartu/commands/register.py
@@ -3,8 +3,8 @@
 import re
 from pathlib import Path
 
-from .command import Command
-from ..utils.registry import Registry
+from urartu.commands.command import Command
+from urartu.utils.registry import Registry
 
 
 @Command.register("register")
diff --git a/urartu/commands/unregister.py b/urartu/commands/unregister.py
@@ -2,8 +2,8 @@
 import logging
 import re
 
-from .command import Command
-from ..utils.registry import Registry
+from urartu.commands.command import Command
+from urartu.utils.registry import Registry
 
 
 @Command.register("unregister")
diff --git a/urartu/config/action_config/default_action_config.yaml b/urartu/config/action_config/default_action_config.yaml
diff --git a/urartu/config/aim/default_aim.yaml b/urartu/config/aim/default_aim.yaml
@@ -0,0 +1,3 @@
+aim:
+  repo: ???
+  log_system_params: true
diff --git a/urartu/config/main.yaml b/urartu/config/main.yaml
@@ -1,42 +1,7 @@
 defaults:
   - _self_
-  - action_config: defaults
+  - action_config: default_action_config
+  - slurm: default_slurm
+  - aim: default_aim
 
 action_name: ???
-
-aim:
-  repo: "aim://0.0.0.0:53800"
-  log_system_params: true
-
-slurm:
-  # Whether or not to run the job on SLURM
-  use_slurm: false
-  # Name of the job on SLURM
-  name: "example"
-  # Comment of the job on SLURM
-  comment: "example job"
-  # Partition of SLURM on which to run the job. This is a required field if using SLURM.
-  partition: ""
-  account: ""
-  # Where the logs produced by the SLURM jobs will be output
-  log_folder: "./slurm_logs"
-  # Maximum number of hours / minutes needed by the job to complete. Above this limit, the job might be pre-empted.
-  time_hours: 1
-  time_minutes: 0
-  # Additional constraints on the hardware of the nodes to allocate (example 'volta' to select a volta GPU)
-  constraint: ""
-  # GB of RAM memory to allocate for each node
-  mem_gb: 100
-  # TCP port on which the workers will synchronize themselves with torch distributed
-  port_id: 40050
-  # Number of CPUs per GPUs to request on the cluster.
-  num_cpu_per_proc: 4
-  # Number of GPUs per node to request on the cluster.
-  num_gpu_per_node: 4
-  # Number of nodes to request on the cluster.
-  num_nodes: 1
-  # Number of processes per node to request on the cluster.
-  num_proc_per_node: 1
-  # Any other parameters for slurm (e.g. account, hint, distribution, etc.,) as dictated by submitit.
-  # Please see https://github.com/facebookincubator/submitit/issues/23#issuecomment-695217824.
-  additional_parameters: {}
diff --git a/urartu/config/slurm/default_slurm.yaml b/urartu/config/slurm/default_slurm.yaml
@@ -0,0 +1,2 @@
+slurm:
+  use_slurm: ???
diff --git a/urartu/datasets/hf/__init__.py b/urartu/datasets/hf/__init__.py
diff --git a/urartu/datasets/hf/dataset_from_dict.py b/urartu/datasets/hf/dataset_from_dict.py
@@ -0,0 +1,16 @@
+from typing import Any, Dict, List
+
+from datasets import Dataset as HFDataset
+
+from urartu.common.dataset import Dataset
+
+
+class DatasetFromDict(Dataset):
+    def __init__(self, cfg: List[Dict[str, Any]]) -> None:
+        super().__init__(cfg)
+
+    def _get_dataset(self):
+        if "data" not in self.cfg:
+            raise TypeError("Argument 'data' is missing")
+
+        self.dataset = HFDataset.from_dict(dict(self.cfg.data))
diff --git a/urartu/datasets/hf/dataset_from_file.py b/urartu/datasets/hf/dataset_from_file.py
@@ -0,0 +1,23 @@
+from typing import Any, Dict, List
+
+from datasets import load_dataset
+from pathlib import Path
+
+from urartu.common.dataset import Dataset
+
+
+class DatasetFromFile(Dataset):
+    def __init__(self, cfg: List[Dict[str, Any]]) -> None:
+        super().__init__(cfg)
+
+    def _get_dataset(self):
+        if "data_files" not in self.cfg:
+            raise TypeError("Argument 'data_files' is missing")
+        file_extension = Path(self.cfg.data_files).suffix
+        if file_extension.startswith(".json"):
+            file_format = "json"
+        elif file_extension == ".txt":
+            file_format = "text"
+        else:
+            raise KeyError(f"Failed to load data file '{file_extension}'")
+        self.dataset = load_dataset(file_format, data_files=self.cfg.data_files)["train"]
diff --git a/urartu/datasets/hf/dataset_from_hub.py b/urartu/datasets/hf/dataset_from_hub.py
@@ -2,16 +2,18 @@
 
 from datasets import load_dataset
 
-from ..common.dataset import Dataset
+from urartu.common.dataset import Dataset
 
 
-class HFDatasets(Dataset):
+class DatasetFromHub(Dataset):
     def __init__(self, cfg: List[Dict[str, Any]]) -> None:
         super().__init__(cfg)
 
     def _get_dataset(self):
+        if "name" not in self.cfg:
+            raise TypeError("Argument 'name' is missing")
         self.dataset = load_dataset(
             self.cfg.name,
-            self.cfg.get("subset"),
+            subset=self.cfg.get("subset"),
             split=self.cfg.get("split"),
         )
diff --git a/urartu/models/model_causal_language.py b/urartu/models/model_causal_language.py
@@ -3,11 +3,11 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from ..common.device import Device
-from ..common.model import Model
+from urartu.common.device import Device
+from urartu.common.model import Model
 
 
-class CausalLMModel(Model):
+class ModelCausalLanguage(Model):
     def __init__(self, cfg) -> None:
         super().__init__(cfg)
 
diff --git a/urartu/models/model_openai.py b/urartu/models/model_openai.py
@@ -5,10 +5,10 @@
 from langchain_openai import AzureChatOpenAI
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from ..common.model import Model
+from urartu.common.model import Model
 
 
-class OpenAIModel(Model):
+class ModelOpenAI(Model):
     def __init__(self, cfg) -> None:
         super().__init__(cfg)
 
diff --git a/urartu/models/model_pipeline.py b/urartu/models/model_pipeline.py
@@ -3,11 +3,11 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-from ..common.device import Device
-from ..common.model import Model
+from urartu.common.device import Device
+from urartu.common.model import Model
 
 
-class PipelineModel(Model):
+class ModelPipeline(Model):
     def __init__(self, cfg) -> None:
         super().__init__(cfg)
 
diff --git a/urartu/utils/hydra_plugin.py b/urartu/utils/hydra_plugin.py
@@ -14,14 +14,9 @@ class UrartuPlugin(SearchPathPlugin):
     def manipulate_search_path(self, search_path: ConfigSearchPath) -> None:
         if registry_paths:
             for registry_path in registry_paths:
-                paths = [
-                    (f"{registry_path}/configs", "file"),
-                    (f"{registry_path}/configs_{current_user}", "file"),
-                    (f"pkg://{registry_path}/configs", "pkg"),
-                    (f"file://{registry_path}/configs_{current_user}", "file")
-                ]
-                for path, type in paths:
-                    search_path.prepend(provider="urartu", path=path)
-                    logger.debug(f"Prepended {type} path: {path}")
+                search_path.append(provider="urartu", path=f"file://{registry_path}/configs_{current_user}")
+                search_path.append(provider="urartu", path=f"pkg://{registry_path}/configs_{current_user}")
+                search_path.append(provider="urartu", path=f"file://{registry_path}/configs")
+                search_path.append(provider="urartu", path=f"pkg://{registry_path}/configs")
         else:
             logger.warning("No registry paths found for UrartuPlugin.")

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+aim:`
	`2`	`+ repo: ???`
	`3`	`+ log_system_params: true`