Repository: allenai/tango (language: python)

Issue:
Fix issue where StepInfo config could be parsed into a Step
<!-- To ensure we can review your pull request promptly please complete this template entirely. -->

<!-- Please reference the issue number here. You can replace "Fixes" with "Closes" if it makes more sense. -->

Changes proposed in this pull request:
<!-- Please list all changes/additions here. -->
- Prevents `construct_arg` method in `from_params.py` to try to parse the argument as a Step, if it is the `config` argument of the `StepInfo` class. 
- This happened when the `config` contained references to upstream steps, which evaluated the `could_be_step` condition to `True`.
- The previous behavior was creating issues with the WandbWorkspace when trying to update and re-parse the StepInfo.

## Before submitting

<!-- Please complete this checklist BEFORE submitting your PR to speed along the review process. -->
- [x] I've read and followed all steps in the [Making a pull request](https://github.com/allenai/tango/blob/main/CONTRIBUTING.md#making-a-pull-request)
    section of the `CONTRIBUTING` docs.
- [x] If this PR fixes a bug, I've added a test that will fail without my fix.


## After submitting

<!-- Please complete this checklist AFTER submitting your PR to speed along the review process. -->
- [ ] All GitHub Actions jobs for my pull request have passed.



Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 83e770e0e..7883efb78 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Made `BeakerExecutor` more robust to connection, timeout, SSL, and other recoverable HTTP errors.
 - Fixed a bug with the internal scheduling logic of the `BeakerExecutor` which
   could delay submitting some steps in parallel.
+- Fixed a bug where creating a `StepInfo` object from params might result in unnecessary imports.
+- Fixed a bug where canceling the Beaker executor might not work properly.
 
 ## [v0.13.0](https://github.com/allenai/tango/releases/tag/v0.13.0) - 2022-09-07
 
diff --git a/tango/common/registrable.py b/tango/common/registrable.py
index e15e0dc85..5feaf8c9c 100644
--- a/tango/common/registrable.py
+++ b/tango/common/registrable.py
@@ -123,15 +123,18 @@ def from_files(cls, still_other_params) -> MyVocabulary:
                     ...  # construct some_params from files
                     return cls(some_params)
         """
+
+        if _cls_is_step(cls) and name == "ref":
+            raise ConfigurationError(
+                "You cannot use the name 'ref' to register a step. This name is reserved."
+            )
+
         registry = Registrable._registry[cls]
 
         def add_subclass_to_registry(subclass: Type[_T]) -> Type[_T]:
             # Add to registry, raise an error if key has already been used.
             if name in registry:
 
-                def fullname(c: type) -> str:
-                    return f"{c.__module__}.{c.__qualname__}"
-
                 already_in_use_for = registry[name][0]
                 if already_in_use_for.__module__ == "__main__":
                     # Sometimes the same class shows up under module.submodule.Class and __main__.Class, and we
@@ -143,15 +146,15 @@ def fullname(c: type) -> str:
                     return already_in_use_for
                 elif exist_ok:
                     message = (
-                        f"Registering {fullname(subclass)} as a {fullname(cls)} under the name {name} overwrites "
-                        f"existing entry {fullname(already_in_use_for)}, which is fine because you said "
-                        "exist_ok=True."
+                        f"Registering {_fullname(subclass)} as a {_fullname(cls)} under the name {name} "
+                        f"overwrites existing entry {_fullname(already_in_use_for)}, which is fine because "
+                        "you said exist_ok=True."
                     )
                     logger.info(message)
                 else:
                     message = (
-                        f"Attempting to register {fullname(subclass)} as a {fullname(cls)} under the name "
-                        f"'{name}' failed. {fullname(already_in_use_for)} is already registered under that name."
+                        f"Attempting to register {_fullname(subclass)} as a {_fullname(cls)} under the name "
+                        f"'{name}' failed. {_fullname(already_in_use_for)} is already registered under that name."
                     )
                     raise ConfigurationError(message)
             registry[name] = (subclass, constructor)
@@ -178,7 +181,11 @@ def search_modules(cls: Type[_RegistrableT], name: str):
         """
         Search for and import modules where ``name`` might be registered.
         """
-        if could_be_class_name(name) or name in Registrable._registry[cls]:
+        if (
+            could_be_class_name(name)
+            or name in Registrable._registry[cls]
+            or (_cls_is_step(cls) and name == "ref")
+        ):
             return None
 
         def try_import(module):
@@ -340,3 +347,14 @@ def _get_suggestion(name: str, available: List[str]) -> Optional[str]:
         if suggestion in available:
             return suggestion
     return None
+
+
+def _fullname(c: type) -> str:
+    return f"{c.__module__}.{c.__qualname__}"
+
+
+def _cls_is_step(c: type) -> bool:
+    # NOTE (epwalsh): importing the actual Step class here would result in a circular
+    # import, even though the import wouldn't be at the top of the module (believe me, I've tried).
+    # So instead we just check the fully qualified name of the class.
+    return _fullname(c) == "tango.step.Step"
diff --git a/tango/integrations/beaker/executor.py b/tango/integrations/beaker/executor.py
index adacfa79c..2a965afb1 100644
--- a/tango/integrations/beaker/executor.py
+++ b/tango/integrations/beaker/executor.py
@@ -18,6 +18,7 @@
     EnvVar,
     ExperimentSpec,
     JobFailedError,
+    JobTimeoutError,
     TaskResources,
     TaskSpec,
     TaskStoppedError,
@@ -391,7 +392,7 @@ def future_done_callback(future: concurrent.futures.Future):
                         )
                     elif isinstance(exc, StepFailedError):
                         failed[step_name] = ExecutionMetadata(logs_location=exc.experiment_url)
-                    elif isinstance(exc, ExecutorError):
+                    elif isinstance(exc, (ExecutorError, CancellationError)):
                         failed[step_name] = ExecutionMetadata()
                     else:
                         log_exception(exc, logger)
@@ -491,6 +492,8 @@ def _execute_sub_graph_for_step(
     ) -> Optional[str]:
         if not in_thread:
             self._is_cancelled.clear()
+        else:
+            self._check_if_cancelled()
 
         step = step_graph[step_name]
 
@@ -521,7 +524,16 @@ def _execute_sub_graph_for_step(
 
         # Follow the experiment and stream the logs until it completes.
         try:
-            self.beaker.experiment.wait_for(experiment, strict=True, quiet=True, poll_interval=2.0)
+            while True:
+                try:
+                    self._check_if_cancelled()
+                    self.beaker.experiment.wait_for(
+                        experiment, strict=True, quiet=True, timeout=2.0
+                    )
+                    time.sleep(2.0)
+                    break
+                except JobTimeoutError:
+                    continue
         except (JobFailedError, TaskStoppedError):
             cli_logger.error(
                 '[red]\N{ballot x} Step [b]"%s"[/] failed. You can check the logs at [b]%s[/][/]',
diff --git a/tango/integrations/beaker/workspace.py b/tango/integrations/beaker/workspace.py
index 3ffda1d76..39095d342 100644
--- a/tango/integrations/beaker/workspace.py
+++ b/tango/integrations/beaker/workspace.py
@@ -1,6 +1,7 @@
 import json
 import random
 import tempfile
+import warnings
 from collections import OrderedDict
 from contextlib import contextmanager
 from pathlib import Path
@@ -118,7 +119,16 @@ def step_starting(self, step: Step) -> None:
         self.locks[step] = lock
 
         step_info = self.step_info(step)
-        if step_info.state not in {StepState.INCOMPLETE, StepState.FAILED, StepState.UNCACHEABLE}:
+        if step_info.state == StepState.RUNNING:
+            # Since we've acquired the step lock we know this step can't be running
+            # elsewhere. But the step state can still say its running if the last
+            # run exited before this workspace had a chance to update the step info.
+            warnings.warn(
+                f"Step info for step '{step.unique_id}' is invalid - says step is running "
+                "although it shouldn't be. Ignoring and overwriting step start time.",
+                UserWarning,
+            )
+        elif step_info.state not in {StepState.INCOMPLETE, StepState.FAILED, StepState.UNCACHEABLE}:
             self.locks.pop(step).release()
             raise StepStateError(
                 step,
@@ -185,6 +195,8 @@ def step_failed(self, step: Step, e: BaseException) -> None:
             self.locks.pop(step).release()
 
     def register_run(self, targets: Iterable[Step], name: Optional[str] = None) -> Run:
+        import concurrent.futures
+
         all_steps = set(targets)
         for step in targets:
             all_steps |= step.recursive_dependencies
@@ -210,23 +222,29 @@ def register_run(self, targets: Iterable[Step], name: Optional[str] = None) -> R
             except DatasetConflict:
                 raise ValueError(f"Run name '{name}' is already in use")
 
-        # Collect step info and add data to run dataset.
         steps: Dict[str, StepInfo] = {}
         run_data: Dict[str, str] = {}
-        with tempfile.TemporaryDirectory() as tmp_dir_name:
-            tmp_dir = Path(tmp_dir_name)
 
+        # Collect step info.
+        with concurrent.futures.ThreadPoolExecutor(
+            thread_name_prefix="BeakerWorkspace.register_run()-"
+        ) as executor:
+            step_info_futures = []
             for step in all_steps:
                 if step.name is None:
                     continue
+                step_info_futures.append(executor.submit(self.step_info, step))
+            for future in concurrent.futures.as_completed(step_info_futures):
+                step_info = future.result()
+                assert step_info.step_name is not None
+                steps[step_info.step_name] = step_info
+                run_data[step_info.step_name] = step_info.unique_id
 
-                step_info = self.step_info(step)
-                steps[step.name] = step_info
-                run_data[step.name] = step.unique_id
-
+        # Create Beaker dataset for run.
+        with tempfile.TemporaryDirectory() as tmp_dir_name:
+            tmp_dir = Path(tmp_dir_name)
             with open(tmp_dir / Constants.RUN_DATA_FNAME, "w") as f:
                 json.dump(run_data, f)
-
             self.beaker.dataset.sync(run_dataset, tmp_dir / Constants.RUN_DATA_FNAME, quiet=True)
 
         return Run(name=cast(str, name), steps=steps, start_date=run_dataset.created)
diff --git a/tango/step_info.py b/tango/step_info.py
index e238b0d45..ac13e2832 100644
--- a/tango/step_info.py
+++ b/tango/step_info.py
@@ -301,7 +301,7 @@ def from_json_dict(cls, json_dict: Dict[str, Any]) -> "StepInfo":
 
         :param json_dict: A dictionary representation, such as the one produced by :meth:`to_json_dict()`.
         """
-        return cls.from_params(
+        step_info = cls.from_params(
             {
                 k: (
                     datetime.strptime(v, "%Y-%m-%dT%H:%M:%S").replace(tzinfo=pytz.utc)
@@ -309,8 +309,11 @@ def from_json_dict(cls, json_dict: Dict[str, Any]) -> "StepInfo":
                     else v
                 )
                 for k, v in json_dict.items()
+                if k != "config"
             }
         )
+        step_info.config = json_dict.get("config")
+        return step_info
 
     @classmethod
     def new_from_step(cls, step: Step, **kwargs) -> "StepInfo":


Test patch:
diff --git a/tests/common/registrable_test.py b/tests/common/registrable_test.py
index 495812d13..463891402 100644
--- a/tests/common/registrable_test.py
+++ b/tests/common/registrable_test.py
@@ -3,6 +3,7 @@
 from tango.common.exceptions import ConfigurationError
 from tango.common.registrable import Registrable
 from tango.common.testing import TangoTestCase
+from tango.step import Step
 
 
 class TestRegistrable(TangoTestCase):
@@ -39,3 +40,10 @@ class MockAlternate2(MockBaseClass):
         with pytest.raises(ConfigurationError) as exc:
             MockBaseClass.by_name("mock_1")
             assert "did you mean 'mock-1'?" in str(exc.value)
+
+    def test_registering_step_by_reserved_name(self):
+        with pytest.raises(ConfigurationError, match="cannot use the name 'ref'"):
+
+            @Step.register("ref")
+            class BadStep(Step):
+                pass