simplify invocations

Author: A9isha

src/maxtext/configs/pyconfig.py

@@ -77,18 +77,22 @@ def _module_from_path(path: str) -> str | None:
   return None
 
 
-def _resolve_or_infer_config(argv: list[str], **kwargs) -> tuple[str, list[str]]:
+def _resolve_or_infer_config(argv: list[str] | None = None, **kwargs) -> tuple[str, list[str]]:
   """Resolves or infers config file path from module."""
+  if argv is None:
+    argv = [""]
   if len(argv) >= 2 and argv[1].endswith(".yml"):
     return resolve_config_path(argv[1]), argv[2:]
-  module = _module_from_path(argv[0])
+  module = _module_from_path(argv[0]) if len(argv) > 0 else None
   if module not in _CONFIG_FILE_MAPPING:
-    raise ValueError(
-        f"No config file provided and no default config found for module '{module}'"
+    config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml")
+    logger.warning(
+        "No config file provided and no default config found for module '%s', using base.yml", module
     )
-  config_path = os.path.join(MAXTEXT_CONFIGS_DIR, _CONFIG_FILE_MAPPING[module])
-  logger.warning("No config file provided, using default config mapping: %s", config_path)
-  remaining_argv = argv[1:]
+  else:
+    config_path = os.path.join(MAXTEXT_CONFIGS_DIR, _CONFIG_FILE_MAPPING[module])
+    logger.warning("No config file provided, using default config mapping: %s", config_path)
+  remaining_argv = argv[1:] if len(argv) > 1 else []
 
   return config_path, remaining_argv
 
@@ -299,14 +303,14 @@ def get_keys(self) -> dict[str, Any]:
     return self._flat_config
 
 
-def initialize(argv: list[str], **kwargs) -> HyperParameters:
+def initialize(argv: list[str] | None = None, **kwargs) -> HyperParameters:
   """Initializes the configuration by loading YAML files, and applying CLI, env, and kwarg overrides."""
   pydantic_config = initialize_pydantic(argv, **kwargs)
   config = HyperParameters(pydantic_config)
   return config
 
 
-def initialize_pydantic(argv: list[str], **kwargs) -> MaxTextConfig:
+def initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConfig:
   """Initializes the configuration by loading YAML files, and applying CLI, env, and kwarg overrides.
   Returns pydantic MaxTextConfig class whereas `initialize` returns the og `HyperParameters`
   """
@@ -446,3 +450,10 @@ def initialize_pydantic(argv: list[str], **kwargs) -> MaxTextConfig:
 # Shim for backward compatibility with pyconfig_deprecated_test.py
 validate_and_update_keys = pyconfig_deprecated.validate_and_update_keys
 __all__ = ["initialize", "initialize_pydantic"]
+
+class _CallablePyconfigModule(sys.modules[__name__].__class__):
+  """Allows calling the module directly as mt.pyconfig()."""
+  def __call__(self, argv: list[str] | None = None, **kwargs) -> HyperParameters:
+    return initialize(argv, **kwargs)
+
+sys.modules[__name__].__class__ = _CallablePyconfigModule

src/maxtext/utils/model_creation_utils.py

@@ -115,9 +115,14 @@ def create_model(config, mesh, model_mode: str = MODEL_MODE_TRAIN, rngs: nnx.Rng
   return model
 
 
-def setup_configs_and_devices(argv: list[str], kwargs):
+def setup_configs_and_devices(argv: list[str] | None = None, kwargs: dict | None = None, **extra_kwargs):
   """Setup device allocation and configs for training and inference."""
-  config = pyconfig.initialize_pydantic(argv, **kwargs)
+  if argv is None:
+    argv = [""]
+  
+  combined_kwargs = dict(kwargs) if kwargs else {}
+  combined_kwargs.update(extra_kwargs)
+  config = pyconfig.initialize_pydantic(argv, **combined_kwargs)
   devices = jax.devices()
   if config.num_trainer_slices == -1 and config.num_samplers_slices == -1:
     max_logging.log("Running on a single slice")
@@ -172,22 +177,24 @@ def setup_configs_and_devices(argv: list[str], kwargs):
         )
       trainer_fsdp = trainer_devices_per_slice // tp
 
-    trainer_update = {
+    trainer_kwargs = dict(combined_kwargs)
+    trainer_kwargs.update({
         "num_slices": config.num_trainer_slices,
         "ici_fsdp_parallelism": trainer_fsdp,
         "ici_tensor_parallelism": tp,
         "dcn_data_parallelism": config.num_trainer_slices,
-    }
+    })
 
-    sampler_update = {
+    sampler_kwargs = dict(combined_kwargs)
+    sampler_kwargs.update({
         "num_slices": config.num_samplers_slices,
         "ici_fsdp_parallelism": len(sampler_devices) // config.num_samplers_slices,
         "ici_tensor_parallelism": -1,
         "dcn_data_parallelism": config.num_samplers_slices,
-    }
+    })
 
-    trainer_config = pyconfig.initialize_pydantic(argv, **trainer_update)
-    sampler_config = pyconfig.initialize_pydantic(argv, **sampler_update)
+    trainer_config = pyconfig.initialize_pydantic(argv, **trainer_kwargs)
+    sampler_config = pyconfig.initialize_pydantic(argv, **sampler_kwargs)
 
   else:
     raise ValueError("num_trainer_slices and num_samplers_slices should be both -1 or positive")