make build argument for nightly builds

Makefile

@@ -7,7 +7,8 @@ DOCDIR := docs
 MDCHECKGLOBS := 'docs/**/*.md' 'docs/**/*.rst' 'examples/**/*.md' 'notebooks/**/*.md' 'scripts/**/*.md'
 MDCHECKFILES := CODE_OF_CONDUCT.md CONTRIBUTING.md DEVELOPING.md README.md
 
-TARGETS := ""  # targets for running pytests: keras,onnx,pytorch,pytorch_models,pytorch_datasets,tensorflow_v1,tensorflow_v1_datasets
+BUILD_ARGS :=  # set nightly to build nightly release
+TARGETS := ""  # targets for running pytests: keras,onnx,pytorch,pytorch_models,pytorch_datasets,tensorflow_v1,tensorflow_v1_models,tensorflow_v1_datasets
 PYTEST_ARGS := ""
 ifneq ($(findstring keras,$(TARGETS)),keras)
     PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/sparseml/keras
@@ -63,7 +64,7 @@ docs:
 
 # creates wheel file
 build:
-	python3 setup.py sdist bdist_wheel
+	python3 setup.py sdist bdist_wheel $(BUILD_ARGS)
 
 # clean package
 clean:

setup.py

@@ -12,11 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
+from datetime import date
 from typing import Dict, List, Tuple
 
 from setuptools import find_packages, setup
 
 
+_PACKAGE_NAME = "sparseml"
+_VERSION = "0.1.0"
+_NIGHTLY = "nightly" in sys.argv
+
+if _NIGHTLY:
+    _PACKAGE_NAME += "-nightly"
+    _VERSION += "." + date.today().strftime("%Y%m%d")
+    # remove nightly param so it does not break bdist_wheel
+    sys.argv.remove("nightly")
+
 _deps = [
     "jupyter>=1.0.0",
     "ipywidgets>=7.0.0",
@@ -32,10 +44,10 @@
     "requests>=2.0.0",
     "scikit-image>=0.15.0",
     "scipy>=1.0.0",
-    "sparsezoo>=0.1.0",
     "tqdm>=4.0.0",
     "toposort>=1.0",
 ]
+_nm_deps = [f"{'sparsezoo-nightly' if _NIGHTLY else 'sparsezoo'}~={_VERSION}"]
 _pytorch_deps = ["torch>=1.1.0", "tensorboard>=1.0", "tensorboardX>=1.0"]
 _pytorch_vision_deps = _pytorch_deps + ["torchvision>=0.3.0"]
 _tensorflow_v1_deps = ["tensorflow<2.0.0", "tensorboard<2.0.0", "tf2onnx>=1.0.0,<1.6"]
@@ -72,7 +84,7 @@ def _setup_package_dir() -> Dict:
 
 
 def _setup_install_requires() -> List:
-    return _deps
+    return _nm_deps + _deps
 
 
 def _setup_extras() -> Dict:
@@ -95,8 +107,8 @@ def _setup_long_description() -> Tuple[str, str]:
 
 
 setup(
-    name="sparseml",
-    version="0.1.0",
+    name=_PACKAGE_NAME,
+    version=_VERSION,
     author="Neuralmagic, Inc.",
     author_email="support@neuralmagic.com",
     description=(

src/sparseml/tensorflow_v1/utils/variable.py

@@ -21,6 +21,7 @@
 try:
     import tensorflow.contrib.graph_editor as graph_editor
     from tensorflow.contrib.graph_editor.util import ListView
+
     tf_contrib_err = None
 except Exception as err:
     graph_editor = None