fix: bug with AIM server availability check

Signed-off-by: Harikrishnan Balagopal <[email protected]>

tuning/aim_loader.py

@@ -12,11 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Standard
-import os
-
 # Local
-from tuning.utils.import_utils import is_aim_available
+from tuning.utils.import_utils import get_aim_config, is_aim_available
 
 if is_aim_available():
     # Third Party
@@ -25,9 +22,7 @@
 
 def get_aimstack_callback():
     # Initialize a new run
-    aim_server = os.environ.get("AIMSTACK_SERVER")
-    aim_db = os.environ.get("AIMSTACK_DB")
-    aim_experiment = os.environ.get("AIMSTACK_EXPERIMENT")
+    aim_server, aim_db, aim_experiment = get_aim_config()
     if aim_experiment is None:
         aim_experiment = ""
 

tuning/utils/import_utils.py

@@ -13,10 +13,27 @@
 # limitations under the License.
 
 # Third Party
+import os
+from typing import Optional, Tuple
 from transformers.utils.import_utils import _is_package_available
 
 _is_aim_available = _is_package_available("aim")
+_aim_server = os.environ.get("AIMSTACK_SERVER")
+_aim_db = os.environ.get("AIMSTACK_DB")
+_aim_experiment = os.environ.get("AIMSTACK_EXPERIMENT")
+
+
+def get_aim_config() -> Tuple[Optional[str], Optional[None], Optional[None]]:
+    """
+    Returns: aim_server, aim_db, aim_experiment
+    """
+    return _aim_server, _aim_db, _aim_experiment
 
 
 def is_aim_available():
-    return _is_aim_available
+    return (
+        _is_aim_available
+        and (_aim_server is not None)
+        and (_aim_db is not None)
+        and (_aim_experiment is not None)
+    )